def first_factors

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

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/pointwise_proof.bend as PointwiseProof
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/autograd/tensor_vjp_LAWS.bend as Laws
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
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
import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as PermuteRecords
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_proof.bend as MatmulProof

301# The pairs the first product g @ y^T reads, on the copy of g and the
302# transposed copy of y, are the pairs right_transposed_factors names on g and
303# y themselves.
304def first_factors(+row: Nat, +column: Nat, +inner: U32, +right_inner: U32, +columns: U32,
305  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
306  +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
307  +plan: Layout.Plan<right, Laws.swapped(right)>, +output_ops: List<&2, StorageLaws.Op>,
308  +moved: {Contracts.words(Laws.transposed(right, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Laws.transposed(right, plan)))) ==
309    PermuteRecords.permuted_words(right, right_fill, right_ops, Laws.swapped(right),
310      Shape.logical_numel(Shape.dimensions(Laws.transposed(right, plan)))) : List<&2, F32>},
311  +re: {Shape.dimensions(Laws.transposed(right, plan)) == [right_inner, columns] : List<&2, U32>},
312  +compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat},
313  +column_below: {Nat.is_lt(column, U32.to_nat(columns)) == True{} : Bool},
314  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
315  {Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, Replayed.snoc(ops, StorageLaws.Second{}),
316    Laws.transposed(right, plan), 0.0, output_ops, count) ==
317    Laws.right_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), shape, fill, ops,
318      right, right_fill, right_ops, Laws.swapped(right), count) : List<&2, DotFold.Factors>}:
319  match count:
320    case 0n:
321      {==}
322    case 1n+previous:
323      +inner_count = U32.to_nat(inner)
324      +column_count = U32.to_nat(columns)
325      +axes = Laws.swapped(right)
326      +output = Laws.transposed(right, plan)
327      +position = Nat.add(Nat.mul(previous, column_count), column)
328      +below = MatmulProof.lt_bound(previous, inner_count, U32.to_nat(right_inner), compatible,
329        Conversion.le_successor_lt(previous, inner_count, room))
330      +bound = MatmulProof.cell_below(output, right_inner, columns, re, previous, below, column, column_below)
331      +cell = U32.from_nat(Nat.add(Nat.mul(row, inner_count), previous))
332      factors_snoc(
333        Contracts.factors(row, column, inner_count, column_count, shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), output, 0.0, output_ops, previous),
334        Laws.right_transposed_factors(row, column, inner_count, column_count, shape, fill, ops, right, right_fill, right_ops, axes, previous),
335        Replayed.value(shape, fill, Replayed.snoc(ops, StorageLaws.Second{}), cell),
336        Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
337        Replayed.value(shape, fill, ops, cell),
338        Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, position))),
339        first_factors(row, column, inner, right_inner, columns, shape, fill, ops, right, right_fill, right_ops, plan, output_ops, moved,
340          re, compatible, column_below, previous, Conversion.le_predecessor(previous, inner_count, room)),
341        SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.Second{}, cell, {==}),
342        Equal.trans(F32, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
343          Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, U32.to_nat(U32.from_nat(position))))),
344          Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, position))),
345          permuted_at(right, right_fill, right_ops, axes, plan, output_ops, moved, U32.from_nat(position),
346            PointwiseProof.word_inside(position, output, bound)),
347          Equal.cong(Nat, F32, number => Replayed.value(right, right_fill, right_ops, U32.from_nat(Values.permute_source_index(right, axes, number))),
348            U32.to_nat(U32.from_nat(position)), position, PointwiseProof.word_exact(position, output, bound))))