def nt_word

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

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

937# Factor k of cell (row, column) of a `Kept` right operand of dimensions
938# [columns, inner] is the word right^T holds at k * columns + column.
939def nt_word(+right: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +columns: U32, +inner: U32,
940  +matrix: {Shape.dimensions(right) == [columns, inner] : List<&2, U32>},
941  +column: Nat, +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
942  +k: Nat, +below: {Nat.is_lt(k, U32.to_nat(inner)) == True{} : Bool}) ->
943  {Replayed.value(right, fill, ops, U32.from_nat(Nat.add(Nat.mul(column, U32.to_nat(inner)), k))) ==
944    Replayed.value(right, fill, ops, U32.from_nat(Values.permute_source_index(right, Laws.swapped(right),
945      Nat.add(Nat.mul(k, U32.to_nat(columns)), column)))) : F32}:
946  +across = U32.to_nat(columns)
947  +span = U32.to_nat(inner)
948  +source = Nat.add(Nat.mul(column, span), k)
949  +extents = Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(right), [columns, inner], matrix)
950  %Equal.sym(List<&2, U32>, Laws.swapped(right), [1, 0], swapped_matrix(right, columns, inner, matrix)) :
951    {Replayed.value(right, fill, ops, U32.from_nat(source)) ==
952      Replayed.value(right, fill, ops, U32.from_nat(Values.permute_source_index(right, _, Nat.add(Nat.mul(k, across), column)))) : F32}
953  %Equal.sym(Maybe<&2, U32>, Values.permute_source(right, [1, 0], U32.from_nat(Nat.add(Nat.mul(k, across), column))), Some{U32.from_nat(source)},
954    TransposeProof.permute_source_transposed(right, columns, inner, matrix, k, below, column, column_below)) :
955    {Replayed.value(right, fill, ops, U32.from_nat(source)) ==
956      Replayed.value(right, fill, ops, U32.from_nat(U32.to_nat(Maybe.default(&2, U32, _, 0)))) : F32}
957  %Equal.sym(Nat, U32.to_nat(U32.from_nat(source)), source,
958    MatrixIndex.word_exact(source, right, MatrixIndex.left_below(across, span, column, column_below, right, extents, k, below))) :
959    {Replayed.value(right, fill, ops, U32.from_nat(source)) == Replayed.value(right, fill, ops, U32.from_nat(_)) : F32}
960  {==}