def tn_word

bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_PROOF.bend:962

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
import bendtorch-core@0.1.0.0/src/tensor/transpose_PROOF.bend as TransposeProof
import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex

962# Factor k of cell (row, column) of a `Swapped` left operand of dimensions
963# [inner, rows] is the word left^T holds at row * inner + k.
964def tn_word(+left: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +inner: U32, +rows: U32,
965  +matrix: {Shape.dimensions(left) == [inner, rows] : List<&2, U32>},
966  +row: Nat, +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
967  +k: Nat, +below: {Nat.is_lt(k, U32.to_nat(inner)) == True{} : Bool}) ->
968  {Replayed.value(left, fill, ops, U32.from_nat(Nat.add(Nat.mul(k, U32.to_nat(rows)), row))) ==
969    Replayed.value(left, fill, ops, U32.from_nat(Values.permute_source_index(left, Laws.swapped(left),
970      Nat.add(Nat.mul(row, U32.to_nat(inner)), k)))) : F32}:
971  +height = U32.to_nat(rows)
972  +span = U32.to_nat(inner)
973  +source = Nat.add(Nat.mul(k, height), row)
974  +extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(left), [inner, rows], matrix)
975  %Equal.sym(List<&2, U32>, Laws.swapped(left), [1, 0], swapped_matrix(left, inner, rows, matrix)) :
976    {Replayed.value(left, fill, ops, U32.from_nat(source)) ==
977      Replayed.value(left, fill, ops, U32.from_nat(Values.permute_source_index(left, _, Nat.add(Nat.mul(row, span), k)))) : F32}
978  %Equal.sym(Maybe<&2, U32>, Values.permute_source(left, [1, 0], U32.from_nat(Nat.add(Nat.mul(row, span), k))), Some{U32.from_nat(source)},
979    TransposeProof.permute_source_transposed(left, inner, rows, matrix, row, row_below, k, below)) :
980    {Replayed.value(left, fill, ops, U32.from_nat(source)) ==
981      Replayed.value(left, fill, ops, U32.from_nat(U32.to_nat(Maybe.default(&2, U32, _, 0)))) : F32}
982  %Equal.sym(Nat, U32.to_nat(U32.from_nat(source)), source,
983    MatrixIndex.word_exact(source, left, MatrixIndex.right_below(span, height, row, row_below, left, extents, k, below))) :
984    {Replayed.value(left, fill, ops, U32.from_nat(source)) == Replayed.value(left, fill, ops, U32.from_nat(_)) : F32}
985  {==}