def pair_contract

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

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/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
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/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
import bendtorch-core@0.1.0.0/src/autograd/tensor_vjp_LAWS.bend as Laws

107# Two delivered contributions, each a replayed tensor with its target's
108# dimensions whose cells hold the stated words, meet the pair contract.
109def pair_contract(+a: Nat, +left: Shape.Shape, -left_cell: U32 -> Maybe<&2, F32>,
110  +b: Nat, +right: Shape.Shape, -right_cell: U32 -> Maybe<&2, F32>, -saved: Backward.Saved, -g: Core.Dynamic,
111  +first: Shape.Shape, +first_dimensions: {Shape.dimensions(first) == Shape.dimensions(left) : List<&2, U32>},
112  +first_fill: F32, +first_ops: List<&2, StorageLaws.Op>,
113  first_values: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
114    {Some{Replayed.value(first, first_fill, first_ops, index)} == left_cell(index) : Maybe<&2, F32>},
115  -first_fits: {Lifecycle.logical_length(Replayed.storage(first, first_fill, first_ops)) == Shape.logical_numel(Shape.dimensions(first)) : Nat},
116  +second: Shape.Shape, +second_dimensions: {Shape.dimensions(second) == Shape.dimensions(right) : List<&2, U32>},
117  +second_fill: F32, +second_ops: List<&2, StorageLaws.Op>,
118  second_values: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
119    {Some{Replayed.value(second, second_fill, second_ops, index)} == right_cell(index) : Maybe<&2, F32>},
120  -second_fits: {Lifecycle.logical_length(Replayed.storage(second, second_fill, second_ops)) == Shape.logical_numel(Shape.dimensions(second)) : Nat}) ->
121  Laws.PairContract<a, left, left_cell, b, right, right_cell,
122    Backward.Pass{saved, g, Done{[Backward.Contribution{a, left, Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}},
123      Backward.Contribution{b, right, Core.Packed{second, Core.TTensor{Replayed.storage(second, second_fill, second_ops), second_fits}}}]}}>:
124  Laws.PairContract{
125    {==},
126    Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(first), Shape.dimensions(left),
127      first_dimensions),
128    index => inside =>
129      +inside = inside
130      Equal.trans(Maybe<&2, F32>,
131        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{first, Core.TTensor{Replayed.storage(first, first_fill, first_ops), first_fits}}, index)),
132        Some{Replayed.value(first, first_fill, first_ops, index)}, left_cell(index),
133        PointwiseProof.replayed_point(first, first_fill, first_ops, first_fits, index,
134          SumToSizeFoldProof.shape_below(first, left, index, first_dimensions, inside)),
135        first_values(index, inside)),
136    {==},
137    Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(second), Shape.dimensions(right),
138      second_dimensions),
139    index => inside =>
140      +inside = inside
141      Equal.trans(Maybe<&2, F32>,
142        Core.dynamic_read_value(Core.read_dynamic(Core.Packed{second, Core.TTensor{Replayed.storage(second, second_fill, second_ops), second_fits}}, index)),
143        Some{Replayed.value(second, second_fill, second_ops, index)}, right_cell(index),
144        PointwiseProof.replayed_point(second, second_fill, second_ops, second_fits, index,
145          SumToSizeFoldProof.shape_below(second, right, index, second_dimensions, inside)),
146        second_values(index, inside))}