def frozen_factors

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

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/tensor/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
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

792# The pairs g^T @ x reads on the Second half of g and the Second half of x
793# are the pairs on g and x themselves.
794def frozen_factors(+row: Nat, +column: Nat, +inner: Nat, +rows: Nat, +columns: Nat,
795  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
796  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +count: Nat) ->
797  {Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
798    shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), count) ==
799    Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
800      shape, fill, ops, left, left_fill, left_ops, count) : List<&2, DotFold.Factors>}:
801  match count:
802    case 0n:
803      {==}
804    case 1n+previous:
805      +g_cell = Matrix.read_index(Matrix.Swapped{}, inner, rows, row, previous)
806      +x_cell = Matrix.read_index(Matrix.Swapped{}, inner, columns, column, previous)
807      factors_snoc(
808        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
809          shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), previous),
810        Contracts.product_factors(row, column, Matrix.Swapped{}, Matrix.Swapped{}, inner, rows, columns,
811          shape, fill, ops, left, left_fill, left_ops, previous),
812        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), g_cell),
813        Replayed.value(left, left_fill, Replayed.snoc(left_ops, StorageLaws.Second{}), x_cell),
814        Replayed.value(shape, fill, ops, g_cell),
815        Replayed.value(left, left_fill, left_ops, x_cell),
816        frozen_factors(row, column, inner, rows, columns, shape, fill, ops, left, left_fill, left_ops, previous),
817        SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.Second{}, g_cell, {==}),
818        SumToSizeFoldProof.half_keeps(left, left_fill, left_ops, StorageLaws.Second{}, x_cell, {==}))