src/shape/broadcast_proofs.bend
bendtorch-core@0.1.0.0 · 235 lines · 19 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3
4def fin_broadcast(+n: Nat, +index: Shape.Fin(n)) -> {Shape.Fin.value(n, index) == Shape.broadcast_index(n, Shape.Fin.value(n, index)) : Nat}:
5 match n:
6 case 0n:
7 match index:
8 case 1n+p:
9 match p:
10 case 0n:
11 match index:
12 case Shape.FZero{}:
13 {==}
14 case Shape.FNext{previous}:
15 match previous:
16 case 1n+q:
17 {==}
18
19def axis_left(+left: Nat, +right: Nat, +output: Nat, +axis: Shape.Axis<left, right, output>, +index: Shape.Fin(output)) -> {Shape.Fin.value(left, Shape.Axis.project_left(left, right, output, axis, index)) == Shape.broadcast_index(left, Shape.Fin.value(output, index)) : Nat}:
20 match axis:
21 case Shape.AxisEqual{left_eq, right_eq}:
22 %h@left_eq : {Shape.Fin.value(_, Shape.Axis.project_left(_, right, output, Shape.AxisEqual{h, right_eq}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
23 fin_broadcast(output, index)
24 case Shape.AxisLeftOne{one, right_eq}:
25 %h@one : {Shape.Fin.value(_, Shape.Axis.project_left(_, right, output, Shape.AxisLeftOne{h, right_eq}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
26 {==}
27 case Shape.AxisRightOne{left_eq, one}:
28 %h@left_eq : {Shape.Fin.value(_, Shape.Axis.project_left(_, right, output, Shape.AxisRightOne{h, one}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
29 fin_broadcast(output, index)
30
31def axis_right(+left: Nat, +right: Nat, +output: Nat, +axis: Shape.Axis<left, right, output>, +index: Shape.Fin(output)) -> {Shape.Fin.value(right, Shape.Axis.project_right(left, right, output, axis, index)) == Shape.broadcast_index(right, Shape.Fin.value(output, index)) : Nat}:
32 match axis:
33 case Shape.AxisEqual{left_eq, right_eq}:
34 %h@right_eq : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisEqual{left_eq, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
35 fin_broadcast(output, index)
36 case Shape.AxisLeftOne{one, right_eq}:
37 %h@right_eq : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisLeftOne{one, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
38 fin_broadcast(output, index)
39 case Shape.AxisRightOne{left_eq, one}:
40 %h@one : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisRightOne{left_eq, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
41 {==}
42
43def aligned_left(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Aligned(left, right, output), +coordinate: Shape.Coordinate(output)) -> {Shape.indices(left, Shape.Aligned.project_left(left, right, output, alignment, coordinate)) == Shape.broadcast_indices(left, Shape.indices(output, coordinate)) : List<&2, Nat>}:
44 match left right output:
45 case Nil{} Nil{} Nil{}:
46 {==}
47 case +lh <> +lt +rh <> +rt +oh <> +ot:
48 Shape.AStep{+axis, +tail} = alignment
49 Shape.CStep{+index, +rest} = coordinate
50 %aligned_left(lt, rt, ot, tail, rest) : {Shape.Fin.value(lh, Shape.Axis.project_left(lh, rh, oh, axis, index)) <> Shape.indices(lt, Shape.Aligned.project_left(lt, rt, ot, tail, rest)) == Shape.broadcast_index(lh, Shape.Fin.value(oh, index)) <> _ : List<&2, Nat>}
51 %axis_left(lh, rh, oh, axis, index) : {Shape.Fin.value(lh, Shape.Axis.project_left(lh, rh, oh, axis, index)) <> Shape.indices(lt, Shape.Aligned.project_left(lt, rt, ot, tail, rest)) == _ <> Shape.indices(lt, Shape.Aligned.project_left(lt, rt, ot, tail, rest)) : List<&2, Nat>}
52 {==}
53 case Nil{} Nil{} oh <> ot:
54 match alignment:
55 case Nil{} rh <> rt Nil{}:
56 match alignment:
57 case lh <> lt Nil{} Nil{}:
58 match alignment:
59 case Nil{} rh <> rt oh <> ot:
60 match alignment:
61 case lh <> lt Nil{} oh <> ot:
62 match alignment:
63 case lh <> lt rh <> rt Nil{}:
64 match alignment:
65
66def aligned_right(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Aligned(left, right, output), +coordinate: Shape.Coordinate(output)) -> {Shape.indices(right, Shape.Aligned.project_right(left, right, output, alignment, coordinate)) == Shape.broadcast_indices(right, Shape.indices(output, coordinate)) : List<&2, Nat>}:
67 match left right output:
68 case Nil{} Nil{} Nil{}:
69 {==}
70 case +lh <> +lt +rh <> +rt +oh <> +ot:
71 Shape.AStep{+axis, +tail} = alignment
72 Shape.CStep{+index, +rest} = coordinate
73 %aligned_right(lt, rt, ot, tail, rest) : {Shape.Fin.value(rh, Shape.Axis.project_right(lh, rh, oh, axis, index)) <> Shape.indices(rt, Shape.Aligned.project_right(lt, rt, ot, tail, rest)) == Shape.broadcast_index(rh, Shape.Fin.value(oh, index)) <> _ : List<&2, Nat>}
74 %axis_right(lh, rh, oh, axis, index) : {Shape.Fin.value(rh, Shape.Axis.project_right(lh, rh, oh, axis, index)) <> Shape.indices(rt, Shape.Aligned.project_right(lt, rt, ot, tail, rest)) == _ <> Shape.indices(rt, Shape.Aligned.project_right(lt, rt, ot, tail, rest)) : List<&2, Nat>}
75 {==}
76 case Nil{} Nil{} oh <> ot:
77 match alignment:
78 case Nil{} rh <> rt Nil{}:
79 match alignment:
80 case lh <> lt Nil{} Nil{}:
81 match alignment:
82 case Nil{} rh <> rt oh <> ot:
83 match alignment:
84 case lh <> lt Nil{} oh <> ot:
85 match alignment:
86 case lh <> lt rh <> rt Nil{}:
87 match alignment:
88
89def indices_length(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {List.length(&2, Nat, dims) == List.length(&2, Nat, Shape.indices(dims, coordinate)) : Nat}:
90 match dims:
91 case Nil{}:
92 {==}
93 case +head <> +tail:
94 Shape.CStep{+index, +rest} = coordinate
95 %indices_length(tail, rest) : {1n+List.length(&2, Nat, tail) == 1n+_ : Nat}
96 {==}
97
98def aligned_length_left(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Aligned(left, right, output)) -> {List.length(&2, Nat, left) == List.length(&2, Nat, output) : Nat}:
99 match left right output:
100 case Nil{} Nil{} Nil{}:
101 {==}
102 case +lh <> +lt +rh <> +rt +oh <> +ot:
103 Shape.AStep{axis, +tail} = alignment
104 %aligned_length_left(lt, rt, ot, tail) : {1n+List.length(&2, Nat, lt) == 1n+_ : Nat}
105 {==}
106 case Nil{} Nil{} oh <> ot:
107 match alignment:
108 case Nil{} rh <> rt Nil{}:
109 match alignment:
110 case lh <> lt Nil{} Nil{}:
111 match alignment:
112 case Nil{} rh <> rt oh <> ot:
113 match alignment:
114 case lh <> lt Nil{} oh <> ot:
115 match alignment:
116 case lh <> lt rh <> rt Nil{}:
117 match alignment:
118
119def aligned_length_right(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Aligned(left, right, output)) -> {List.length(&2, Nat, right) == List.length(&2, Nat, output) : Nat}:
120 match left right output:
121 case Nil{} Nil{} Nil{}:
122 {==}
123 case +lh <> +lt +rh <> +rt +oh <> +ot:
124 Shape.AStep{axis, +tail} = alignment
125 %aligned_length_right(lt, rt, ot, tail) : {1n+List.length(&2, Nat, rt) == 1n+_ : Nat}
126 {==}
127 case Nil{} Nil{} oh <> ot:
128 match alignment:
129 case Nil{} rh <> rt Nil{}:
130 match alignment:
131 case lh <> lt Nil{} Nil{}:
132 match alignment:
133 case Nil{} rh <> rt oh <> ot:
134 match alignment:
135 case lh <> lt Nil{} oh <> ot:
136 match alignment:
137 case lh <> lt rh <> rt Nil{}:
138 match alignment:
139
140def pad_length(+count: Nat, +dims: List<&2, Nat>) -> {Nat.add(count, List.length(&2, Nat, dims)) == List.length(&2, Nat, Shape.pad(count, dims)) : Nat}:
141 match count:
142 case 0n:
143 {==}
144 case 1n+rest:
145 %pad_length(rest, dims) : {1n+Nat.add(rest, List.length(&2, Nat, dims)) == 1n+_ : Nat}
146 {==}
147
148def add_succ(a: Nat, -b: Nat) -> {1n+Nat.add(a, b) == Nat.add(a, 1n+b) : Nat}:
149 match a:
150 case 0n:
151 {==}
152 case 1n+p:
153 %add_succ(p, b) : {2n+Nat.add(p, b) == 1n+_ : Nat}
154 {==}
155
156def add_zero(a: Nat) -> {Nat.add(a, 0n) == a : Nat}:
157 match a:
158 case 0n:
159 {==}
160 case 1n+p:
161 %add_zero(p) : {1n+Nat.add(p, 0n) == 1n+_ : Nat}
162 {==}
163
164def sub_add_zero(a: Nat) -> {a == Nat.sub(Nat.add(a, 0n), 0n) : Nat}:
165 match a:
166 case 0n:
167 {==}
168 case 1n+p:
169 %add_zero(p) : {1n+_ == 1n+Nat.add(p, 0n) : Nat}
170 {==}
171
172def sub_add(+a: Nat, +b: Nat) -> {a == Nat.sub(Nat.add(a, b), b) : Nat}:
173 match b:
174 case 0n:
175 sub_add_zero(a)
176 case 1n+q:
177 %add_succ(a, q) : {a == Nat.sub(_, 1n+q) : Nat}
178 sub_add(a, q)
179
180def drop_zero(+values: List<&2, Nat>) -> {values == List.drop(&2, Nat, values, 0n) : List<&2, Nat>}:
181 match values:
182 case Nil{}:
183 {==}
184 case head <> tail:
185 {==}
186
187def broadcast_nil(+dims: List<&2, Nat>) -> {Nil{} == Shape.broadcast_indices(dims, Nil{}) : List<&2, Nat>}:
188 match dims:
189 case Nil{}:
190 {==}
191 case head <> tail:
192 {==}
193
194def drop_broadcast(+count: Nat, +dims: List<&2, Nat>, +values: List<&2, Nat>) -> {List.drop(&2, Nat, Shape.broadcast_indices(Shape.pad(count, dims), values), count) == Shape.broadcast_indices(dims, List.drop(&2, Nat, values, count)) : List<&2, Nat>}:
195 match count:
196 case 0n:
197 %drop_zero(values) : {List.drop(&2, Nat, Shape.broadcast_indices(dims, values), 0n) == Shape.broadcast_indices(dims, _) : List<&2, Nat>}
198 %drop_zero(Shape.broadcast_indices(dims, values)) : {_ == Shape.broadcast_indices(dims, values) : List<&2, Nat>}
199 {==}
200 case 1n+rest:
201 match values:
202 case Nil{}:
203 broadcast_nil(dims)
204 case +head <> +tail:
205 drop_broadcast(rest, dims, tail)
206
207def unpad_indices(+count: Nat, +dims: List<&2, Nat>, +coordinate: Shape.Coordinate(Shape.pad(count, dims))) -> {Shape.indices(dims, Shape.unpad(count, dims, coordinate)) == List.drop(&2, Nat, Shape.indices(Shape.pad(count, dims), coordinate), count) : List<&2, Nat>}:
208 match count:
209 case 0n:
210 drop_zero(Shape.indices(dims, coordinate))
211 case 1n+rest:
212 Shape.CStep{index, +tail} = coordinate
213 unpad_indices(rest, dims, tail)
214
215def project_left(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, +coordinate: Shape.Coordinate(output)) -> {Shape.indices(left, Shape.project_left(left, right, output, alignment, coordinate)) == Shape.project_indices(left, Shape.indices(output, coordinate)) : List<&2, Nat>}:
216 match alignment:
217 case Shape.BAlign{+lp, +rp, minimal, +aligned}:
218 %indices_length(output, coordinate) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(left, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, left)))) : List<&2, Nat>}
219 %aligned_length_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(left, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, left)))) : List<&2, Nat>}
220 %pad_length(lp, left) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(left, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, left)))) : List<&2, Nat>}
221 %sub_add(lp, List.length(&2, Nat, left)) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(left, List.drop(&2, Nat, Shape.indices(output, coordinate), _)) : List<&2, Nat>}
222 %drop_broadcast(lp, left, Shape.indices(output, coordinate)) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == _ : List<&2, Nat>}
223 %aligned_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate) : {Shape.indices(left, Shape.unpad(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == List.drop(&2, Nat, _, lp) : List<&2, Nat>}
224 unpad_indices(lp, left, Shape.Aligned.project_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))
225
226def project_right(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, +coordinate: Shape.Coordinate(output)) -> {Shape.indices(right, Shape.project_right(left, right, output, alignment, coordinate)) == Shape.project_indices(right, Shape.indices(output, coordinate)) : List<&2, Nat>}:
227 match alignment:
228 case Shape.BAlign{+lp, +rp, minimal, +aligned}:
229 %indices_length(output, coordinate) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(right, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, right)))) : List<&2, Nat>}
230 %aligned_length_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(right, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, right)))) : List<&2, Nat>}
231 %pad_length(rp, right) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(right, List.drop(&2, Nat, Shape.indices(output, coordinate), Nat.sub(_, List.length(&2, Nat, right)))) : List<&2, Nat>}
232 %sub_add(rp, List.length(&2, Nat, right)) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == Shape.broadcast_indices(right, List.drop(&2, Nat, Shape.indices(output, coordinate), _)) : List<&2, Nat>}
233 %drop_broadcast(rp, right, Shape.indices(output, coordinate)) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == _ : List<&2, Nat>}
234 %aligned_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate) : {Shape.indices(right, Shape.unpad(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))) == List.drop(&2, Nat, _, rp) : List<&2, Nat>}
235 unpad_indices(rp, right, Shape.Aligned.project_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned, coordinate))