src/autograd/broadcast_pointwise_PROOF.bend
bendtorch-core@0.1.0.0 · 196 lines · 8 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
7import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
8import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
9import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
10import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts
11import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
12import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
13import bendtorch-core@0.1.0.0/src/tensor/contracts/broadcast.bend as PointwiseBroadcast
14import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_LAWS.bend as BinaryLaws
15import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_proofs.bend as Binary
16import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_PROOF.bend as BinaryProof
17import bendtorch-core@0.1.0.0/src/autograd/broadcast.bend as Broadcast
18import bendtorch-core@0.1.0.0/src/autograd/reduce_LAWS.bend as ReduceLaws
19import bendtorch-core@0.1.0.0/src/autograd/reduce_fold_PROOF.bend as FoldProof
20import bendtorch-core@0.1.0.0/src/autograd/reduce_to_fold_PROOF.bend as ReduceTo
21import bendtorch-core@0.1.0.0/src/autograd/broadcast_pointwise_LAWS.bend as Laws
22
23
24
25def cells_keep(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +target: Shape.Shape,
26 +count: Nat) ->
27 {BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, count) ==
28 Laws.cells(plan, shape, fill, ops, count) : List<&2, F32>}:
29 match count:
30 case 0n:
31 {==}
32 case 1n+previous:
33 Equal.cong(List<&2, F32>, List<&2, F32>,
34 before => List.append(&2, F32, before, [Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))]),
35 BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, previous), Laws.cells(plan, shape, fill, ops, previous),
36 cells_keep(plan, shape, fill, ops, target, previous))
37
38def cells_length(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) ->
39 {List.length(&2, F32, Laws.cells(plan, shape, fill, ops, count)) == count : Nat}:
40 match count:
41 case 0n:
42 {==}
43 case 1n+previous:
44 +before = Laws.cells(plan, shape, fill, ops, previous)
45 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))
46 Equal.trans(Nat, List.length(&2, F32, List.append(&2, F32, before, [last])), 1n+List.length(&2, F32, before), 1n+previous,
47 FoldProof.length_snoc(before, last),
48 Equal.cong(Nat, Nat, n => 1n+n, List.length(&2, F32, before), previous, cells_length(plan, shape, fill, ops, previous)))
49
50
51def cells_step(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +previous: Nat, +position: Nat,
52 below: {Nat.is_lt(position, 1n+previous) == True{} : Bool}, decision: Shape.BoolCase<Nat.is_eq(position, previous)>,
53 earlier: @inside: {Nat.is_lt(position, previous) == True{} : Bool} ->
54 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, previous), position) ==
55 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}) ->
56 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, 1n+previous), position) ==
57 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}:
58 match decision:
59 case Shape.BFalse{equal}:
60 +before = Laws.cells(plan, shape, fill, ops, previous)
61 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))
62 +inside = PointwiseProof.lt_succ_ne(position, previous, below, equal)
63 Equal.trans(Maybe<&2, F32>, ReduceLaws.word_at(List.append(&2, F32, before, [last]), position), ReduceLaws.word_at(before, position),
64 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))},
65 FoldProof.word_at_front(before, last, position,
66 %Equal.sym(Nat, List.length(&2, F32, before), previous, cells_length(plan, shape, fill, ops, previous)) :
67 {Nat.is_lt(position, _) == True{} : Bool}
68 inside),
69 earlier(inside))
70 case Shape.BTrue{equal}:
71 +before = Laws.cells(plan, shape, fill, ops, previous)
72 +last = Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, previous)))
73 %Equal.sym(Nat, position, previous, Arith.eq_true(position, previous, Equal.sym(Bool, True{}, Nat.is_eq(position, previous), equal))) :
74 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) ==
75 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, _)))} : Maybe<&2, F32>}
76 %cells_length(plan, shape, fill, ops, previous) :
77 {ReduceLaws.word_at(List.append(&2, F32, before, [last]), _) == Some{last} : Maybe<&2, F32>}
78 FoldProof.word_at_last(before, last)
79
80def cells_at(+plan: Shape.BroadcastPlan, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +position: Nat,
81 below: {Nat.is_lt(position, count) == True{} : Bool}) ->
82 {ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, count), position) ==
83 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>}:
84 match count:
85 case 0n:
86 Empty.absurd({ReduceLaws.word_at(Nil{}, position) ==
87 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, position)))} : Maybe<&2, F32>},
88 Arith.lt_zero(position, below))
89 case 1n+previous:
90 cells_step(plan, shape, fill, ops, previous, position, below, Shape.bool_case(Nat.is_eq(position, previous)),
91 inside => cells_at(plan, shape, fill, ops, previous, position, inside))
92
93
94
95def listed_value(+shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat, +values: List<&2, F32>, +index: U32,
96 words: {Contracts.words(shape, fill, ops, count) == values : List<&2, F32>},
97 below: {Nat.is_lt(U32.to_nat(index), count) == True{} : Bool}) ->
98 {Some{Replayed.value(shape, fill, ops, index)} == ReduceLaws.word_at(values, U32.to_nat(index)) : Maybe<&2, F32>}:
99 %Conversion.word_natural_exact(index) :
100 {Some{Replayed.value(shape, fill, ops, _)} == ReduceLaws.word_at(values, U32.to_nat(index)) : Maybe<&2, F32>}
101 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(shape, fill, ops, U32.from_nat(U32.to_nat(index)))},
102 ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), U32.to_nat(index)), ReduceLaws.word_at(values, U32.to_nat(index)),
103 Equal.sym(Maybe<&2, F32>, ReduceLaws.word_at(Contracts.words(shape, fill, ops, count), U32.to_nat(index)),
104 Some{Replayed.value(shape, fill, ops, U32.from_nat(U32.to_nat(index)))},
105 FoldProof.words_at(shape, fill, ops, count, U32.to_nat(index), below)),
106 Equal.cong(List<&2, F32>, Maybe<&2, F32>, listed => ReduceLaws.word_at(listed, U32.to_nat(index)),
107 Contracts.words(shape, fill, ops, count), values, words))
108
109
110
111
112
113def run_lift(+shape: Shape.Shape, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
114 +planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
115 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
116 +fill: F32, +ops: List<&2, StorageLaws.Op>,
117 -P: Core.Dynamic & Result<&2, &1, Broadcast.Error, Core.Dynamic> -> Type,
118 h: Laws.BroadcastWith(plan, shape, fill, ops, P)) ->
119 @-fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat} ->
120 P(Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)):
121 +output = Shape.output_shape(plan)
122 +count = Shape.logical_numel(Shape.dimensions(output))
123 fits =>
124 -input = {Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}} : Core.Dynamic}
125 %Equal.sym(Result<&2, &2, Shape.Error, Shape.BroadcastPlan>, Shape.broadcast_checked(shape, target), Done{plan}, planned) :
126 P(Broadcast.planned(input, target, _))
127 %Equal.sym(Bool, Layout.words_same(Shape.dimensions(output), Shape.dimensions(target)), True{}, ReduceTo.fitted_same(target, plan, same)) :
128 P(Broadcast.fitted(input, target, plan, _))
129 Replayed.new_lift(target, 0.0, result => P(Broadcast.zeroed(input, target, plan, result)),
130 zeros_fits => Binary.allocated_replayed(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, planned,
131 pair => P(Broadcast.finish(pair)),
132 kept => kept_words => input_fits => zeros_kept => zeros_words => zeros_kept_fits => output_ops => values => output_fits =>
133 h(kept, kept_words, input_fits, output_ops,
134 Equal.trans(List<&2, F32>, Contracts.words(output, 0.0, output_ops, count),
135 BinaryLaws.cells(Elementwise.KeepLeft{}, plan, shape, fill, ops, target, 0.0, Nil{}, count), Laws.cells(plan, shape, fill, ops, count),
136 values, cells_keep(plan, shape, fill, ops, target, count)),
137 output_fits))(fits, zeros_fits))
138
139def output_below(+plan: Shape.BroadcastPlan, +target: Shape.Shape, +index: U32,
140 same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
141 inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool}) ->
142 {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan)))) == True{} : Bool}:
143 %Equal.sym(List<&2, U32>, Shape.dimensions(Shape.output_shape(plan)), Shape.dimensions(target), same) :
144 {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(_)) == True{} : Bool}
145 inside
146
147def broadcast_to_pointwise(+shape: Shape.Shape, +target: Shape.Shape, +plan: Shape.BroadcastPlan,
148 +planned: {Shape.broadcast_checked(shape, target) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>},
149 +same: {Shape.dimensions(Shape.output_shape(plan)) == Shape.dimensions(target) : List<&2, U32>},
150 +fill: F32, +ops: List<&2, StorageLaws.Op>,
151 -fits: {Lifecycle.logical_length(Replayed.storage(shape, fill, ops)) == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) ->
152 Laws.BroadcastToContract<target, shape,
153 index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
154 other => Replayed.value(shape, fill, ops, other),
155 Broadcast.broadcast_to(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, ops), fits}}, target)>:
156 +output = Shape.output_shape(plan)
157 +count = Shape.logical_numel(Shape.dimensions(output))
158 +total = Shape.logical_numel(Shape.dimensions(shape))
159 run_lift(shape, target, plan, planned, same, fill, ops,
160 result => Laws.BroadcastToContract<target, shape,
161 index => Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, index))),
162 other => Replayed.value(shape, fill, ops, other), result>,
163 kept => kept_words => input_fits => output_ops => values => output_fits =>
164 Laws.BroadcastToContract{
165 {==},
166 Equal.cong(List<&2, U32>, Maybe<&2, List<&2, U32>>, dimensions => Some{dimensions}, Shape.dimensions(output), Shape.dimensions(target), same),
167 index => inside =>
168 +below = output_below(plan, target, index, same, inside)
169 Equal.trans(Maybe<&2, F32>,
170 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{output, Core.TTensor{Replayed.storage(output, 0.0, output_ops), output_fits}}, index)),
171 Some{Replayed.value(output, 0.0, output_ops, index)},
172 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, U32.to_nat(index))))},
173 PointwiseProof.replayed_point(output, 0.0, output_ops, output_fits, index, below),
174 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(output, 0.0, output_ops, index)},
175 ReduceLaws.word_at(Laws.cells(plan, shape, fill, ops, count), U32.to_nat(index)),
176 Some{Replayed.value(shape, fill, ops, U32.from_nat(PointwiseBroadcast.left_map(plan, U32.to_nat(index))))},
177 listed_value(output, 0.0, output_ops, count, Laws.cells(plan, shape, fill, ops, count), index, values, below),
178 cells_at(plan, shape, fill, ops, count, U32.to_nat(index), below))),
179 other => inside =>
180 +within = {inside : {Nat.is_lt(U32.to_nat(other), total) == True{} : Bool}}
181 Equal.trans(Maybe<&2, F32>,
182 Core.dynamic_read_value(Core.read_dynamic(Core.Packed{shape, Core.TTensor{Replayed.storage(shape, fill, kept), input_fits}}, other)),
183 Some{Replayed.value(shape, fill, kept, other)}, Some{Replayed.value(shape, fill, ops, other)},
184 PointwiseProof.replayed_point(shape, fill, kept, input_fits, other, within),
185 Equal.trans(Maybe<&2, F32>, Some{Replayed.value(shape, fill, kept, other)},
186 ReduceLaws.word_at(Contracts.words(shape, fill, ops, total), U32.to_nat(other)), Some{Replayed.value(shape, fill, ops, other)},
187 listed_value(shape, fill, kept, total, Contracts.words(shape, fill, ops, total), other, kept_words, within),
188 Equal.sym(Maybe<&2, F32>, Some{Replayed.value(shape, fill, ops, other)},
189 ReduceLaws.word_at(Contracts.words(shape, fill, ops, total), U32.to_nat(other)),
190 listed_value(shape, fill, ops, total, Contracts.words(shape, fill, ops, total), other, {==}, within))))})(fits)
191
192def Laws.broadcast_to_pointwise(shape, target, plan, planned, same, fill, ops, fits):
193 broadcast_to_pointwise(shape, target, plan, planned, same, fill, ops, fits)
194
195def Laws.broadcast_to_replayed(shape, target, plan, planned, same, fill, ops, fits, P, h):
196 run_lift(shape, target, plan, planned, same, fill, ops, P, h)(fits)