def second_factors

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

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

350# The pairs the second product x^T @ g reads, on the transposed copy of x and
351# the copy of g's first half, are the pairs left_transposed_factors names on x
352# and g themselves.
353def second_factors(+row: Nat, +column: Nat, +rows: U32, +inner: U32, +columns: U32,
354  +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
355  +plan: Layout.Plan<left, Laws.swapped(left)>, +output_ops: List<&2, StorageLaws.Op>,
356  +moved: {Contracts.words(Laws.transposed(left, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Laws.transposed(left, plan)))) ==
357    PermuteRecords.permuted_words(left, left_fill, left_ops, Laws.swapped(left),
358      Shape.logical_numel(Shape.dimensions(Laws.transposed(left, plan)))) : List<&2, F32>},
359  +le: {Shape.dimensions(Laws.transposed(left, plan)) == [rows, inner] : List<&2, U32>},
360  +row_below: {Nat.is_lt(row, U32.to_nat(rows)) == True{} : Bool},
361  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>,
362  +count: Nat, +room: {Nat.is_le(count, U32.to_nat(inner)) == True{} : Bool}) ->
363  {Contracts.factors(row, column, U32.to_nat(inner), U32.to_nat(columns), Laws.transposed(left, plan), 0.0, output_ops,
364    shape, fill, Replayed.snoc(Replayed.snoc(ops, StorageLaws.First{}), StorageLaws.Second{}), count) ==
365    Laws.left_transposed_factors(row, column, U32.to_nat(inner), U32.to_nat(columns), left, left_fill, left_ops, Laws.swapped(left),
366      shape, fill, ops, count) : List<&2, DotFold.Factors>}:
367  match count:
368    case 0n:
369      {==}
370    case 1n+previous:
371      +inner_count = U32.to_nat(inner)
372      +column_count = U32.to_nat(columns)
373      +axes = Laws.swapped(left)
374      +output = Laws.transposed(left, plan)
375      +position = Nat.add(Nat.mul(row, inner_count), previous)
376      +bound = MatmulProof.cell_below(output, rows, inner, le, row, row_below, previous,
377        Conversion.le_successor_lt(previous, inner_count, room))
378      +cell = U32.from_nat(Nat.add(Nat.mul(previous, column_count), column))
379      +front = Replayed.snoc(ops, StorageLaws.First{})
380      factors_snoc(
381        Contracts.factors(row, column, inner_count, column_count, output, 0.0, output_ops, shape, fill, Replayed.snoc(front, StorageLaws.Second{}), previous),
382        Laws.left_transposed_factors(row, column, inner_count, column_count, left, left_fill, left_ops, axes, shape, fill, ops, previous),
383        Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
384        Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), cell),
385        Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, position))),
386        Replayed.value(shape, fill, ops, cell),
387        second_factors(row, column, rows, inner, columns, left, left_fill, left_ops, plan, output_ops, moved, le, row_below, shape, fill, ops,
388          previous, Conversion.le_predecessor(previous, inner_count, room)),
389        Equal.trans(F32, Replayed.value(output, 0.0, output_ops, U32.from_nat(position)),
390          Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, U32.to_nat(U32.from_nat(position))))),
391          Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, position))),
392          permuted_at(left, left_fill, left_ops, axes, plan, output_ops, moved, U32.from_nat(position),
393            PointwiseProof.word_inside(position, output, bound)),
394          Equal.cong(Nat, F32, number => Replayed.value(left, left_fill, left_ops, U32.from_nat(Values.permute_source_index(left, axes, number))),
395            U32.to_nat(U32.from_nat(position)), position, PointwiseProof.word_exact(position, output, bound))),
396        Equal.trans(F32, Replayed.value(shape, fill, Replayed.snoc(front, StorageLaws.Second{}), cell), Replayed.value(shape, fill, front, cell),
397          Replayed.value(shape, fill, ops, cell),
398          SumToSizeFoldProof.half_keeps(shape, fill, front, StorageLaws.Second{}, cell, {==}),
399          SumToSizeFoldProof.half_keeps(shape, fill, ops, StorageLaws.First{}, cell, {==})))