def tn_factors

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

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/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values

1018def tn_factors(+row: Nat, +column: Nat, +columns: Nat, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
1019  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +inner: U32, +rows: U32,
1020  +matrix: {Shape.dimensions(left) == [inner, rows] : List<&2, U32>},
1021  +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
1022  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
1023  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, U32.to_nat(inner), U32.to_nat(rows), columns,
1024    left, left_fill, left_ops, right, right_fill, right_ops, count) ==
1025    Laws.left_transposed_factors(row, column, U32.to_nat(inner), columns, left, left_fill, left_ops, Laws.swapped(left),
1026      right, right_fill, right_ops, count) : List<&2, DotFold.Factors>}:
1027  match count:
1028    case 0n:
1029      {==}
1030    case 1n+previous:
1031      +span = U32.to_nat(inner)
1032      +height = U32.to_nat(rows)
1033      factors_snoc(
1034        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, span, height, columns,
1035          left, left_fill, left_ops, right, right_fill, right_ops, previous),
1036        Laws.left_transposed_factors(row, column, span, columns, left, left_fill, left_ops, Laws.swapped(left),
1037          right, right_fill, right_ops, previous),
1038        Replayed.value(left, left_fill, left_ops, U32.from_nat(Nat.add(Nat.mul(previous, height), row))),
1039        Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(previous, columns), column))),
1040        Replayed.value(left, left_fill, left_ops,
1041          U32.from_nat(Values.permute_source_index(left, Laws.swapped(left), Nat.add(Nat.mul(row, span), previous)))),
1042        Replayed.value(right, right_fill, right_ops, U32.from_nat(Nat.add(Nat.mul(previous, columns), column))),
1043        tn_factors(row, column, columns, left, left_fill, left_ops, right, right_fill, right_ops, inner, rows, matrix,
1044          row_below, previous, Conversion.le_predecessor(previous, span, room)),
1045        tn_word(left, left_fill, left_ops, inner, rows, matrix, row, row_below, previous,
1046          Conversion.le_successor_lt(previous, span, room)),
1047        {==})