src/shape/coordinate_proofs.bend
bendtorch-core@0.1.0.0 · 211 lines · 22 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
5import bendtorch-core@0.1.0.0/src/shape/broadcast_proofs.bend as Broadcast
6
7type Fits.Step<-dh: Nat, -dt: List<&2, Nat>, -vh: Nat, -vt: List<&2, Nat>> is Data:
8 FStep{bound: {True{} == Nat.is_lt(vh, dh) : Bool}, tail: Fits(dt, vt)}
9
10def Fits(dims: List<&2, Nat>, values: List<&2, Nat>) -> Data:
11 match dims values:
12 case Nil{} Nil{}:
13 Unit
14 case dh <> dt vh <> vt:
15 Fits.Step<dh, dt, vh, vt>
16 case Nil{} vh <> vt:
17 Empty
18 case dh <> dt Nil{}:
19 Empty
20
21def fin_from(+bound: Nat, +value: Nat, fits: {True{} == Nat.is_lt(value, bound) : Bool}) -> Shape.Fin(bound):
22 match bound value:
23 case 0n 0n:
24 Empty.absurd(Shape.Fin(0n), Shape.clash(fits))
25 case 0n 1n+q:
26 Empty.absurd(Shape.Fin(0n), Shape.clash(fits))
27 case 1n+p 0n:
28 Shape.FZero{}
29 case 1n+p 1n+q:
30 Shape.FNext{fin_from(p, q, fits)}
31
32def fin_from_value(+bound: Nat, +value: Nat, fits: {True{} == Nat.is_lt(value, bound) : Bool}) -> {Shape.Fin.value(bound, fin_from(bound, value, fits)) == value : Nat}:
33 match bound value:
34 case 0n 0n:
35 Empty.absurd({Shape.Fin.value(0n, fin_from(0n, 0n, fits)) == 0n : Nat}, Shape.clash(fits))
36 case 0n 1n+q:
37 Empty.absurd({Shape.Fin.value(0n, fin_from(0n, 1n+q, fits)) == 1n+q : Nat}, Shape.clash(fits))
38 case 1n+p 0n:
39 {==}
40 case 1n+p 1n+q:
41 Equal.cong(Nat, Nat, value => 1n+value, Shape.Fin.value(p, fin_from(p, q, fits)), q, fin_from_value(p, q, fits))
42
43def coordinate_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, fits: Fits(dims, values)) -> Shape.Coordinate(dims):
44 match dims values:
45 case Nil{} Nil{}:
46 Unit{}
47 case +dh <> +dt +vh <> +vt:
48 FStep{bound, tail} = fits
49 Shape.CStep{fin_from(dh, vh, bound), coordinate_of(dt, vt, tail)}
50 case Nil{} vh <> vt:
51 match fits:
52 case dh <> dt Nil{}:
53 match fits:
54
55def indices_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, +fits: Fits(dims, values)) -> {Shape.indices(dims, coordinate_of(dims, values, fits)) == values : List<&2, Nat>}:
56 match dims values:
57 case Nil{} Nil{}:
58 {==}
59 case +dh <> +dt +vh <> +vt:
60 FStep{+bound, +tail} = fits
61 %Equal.sym(Nat, Shape.Fin.value(dh, fin_from(dh, vh, bound)), vh, fin_from_value(dh, vh, bound)) : {_ <> Shape.indices(dt, coordinate_of(dt, vt, tail)) == vh <> vt : List<&2, Nat>}
62 %Equal.sym(List<&2, Nat>, Shape.indices(dt, coordinate_of(dt, vt, tail)), vt, indices_of(dt, vt, tail)) : {vh <> _ == vh <> vt : List<&2, Nat>}
63 {==}
64 case Nil{} vh <> vt:
65 match fits:
66 case dh <> dt Nil{}:
67 match fits:
68
69def offset_indices(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.offset(dims, Shape.indices(dims, coordinate)) == Shape.row_major(dims, coordinate) : Nat}:
70 match dims:
71 case Nil{}:
72 {==}
73 case +head <> +tail:
74 Shape.CStep{+index, +rest} = coordinate
75 Equal.cong(Nat, Nat, value => Nat.add(Nat.mul(Shape.Fin.value(head, index), Shape.numel(tail)), value), Shape.offset(tail, Shape.indices(tail, rest)), Shape.row_major(tail, rest), offset_indices(tail, rest))
76
77def row_major_lt(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> Shape.LT(Shape.row_major(dims, coordinate), Shape.numel(dims)):
78 %Index.flatten_row_major(dims, coordinate) : Shape.LT(_, Shape.numel(dims))
79 Shape.Fin.bound(Shape.numel(dims), Shape.flatten(dims, coordinate))
80
81def size_fits(+shape: Shape.Shape, -index: Nat, fits: {True{} == Nat.is_lt(index, Shape.elements(shape)) : Bool}) -> {True{} == Nat.is_lt(index, Shape.numel(Shape.extents(shape))) : Bool}:
82 %Shape.elements_exact(shape) : {True{} == Nat.is_lt(index, _) : Bool}
83 fits
84
85def size_misses(+shape: Shape.Shape, -index: Nat, misses: {False{} == Nat.is_lt(index, Shape.elements(shape)) : Bool}) -> {False{} == Nat.is_lt(index, Shape.numel(Shape.extents(shape))) : Bool}:
86 %Shape.elements_exact(shape) : {False{} == Nat.is_lt(index, _) : Bool}
87 misses
88
89def offset_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, +fits: Fits(dims, values)) -> {Shape.row_major(dims, coordinate_of(dims, values, fits)) == Shape.offset(dims, values) : Nat}:
90 %indices_of(dims, values, fits) : {Shape.row_major(dims, coordinate_of(dims, values, fits)) == Shape.offset(dims, _) : Nat}
91 Equal.sym(Nat, Shape.offset(dims, Shape.indices(dims, coordinate_of(dims, values, fits))), Shape.row_major(dims, coordinate_of(dims, values, fits)), offset_indices(dims, coordinate_of(dims, values, fits)))
92
93def digits_step(+head: Nat, +tail: List<&2, Nat>, rec: @+index: Shape.Fin(Shape.numel(tail)) -> {Shape.indices(tail, Shape.unflatten(tail, index)) == Shape.digits(tail, Shape.Fin.value(Shape.numel(tail), index)) : List<&2, Nat>}, coordinate: Shape.Fin(head) & Shape.Fin(Shape.numel(tail))) -> {Shape.indices(head <> tail, Shape.unflatten.step(head, tail, rest => Shape.unflatten(tail, rest), coordinate)) == Shape.digits(head <> tail, Shape.Fin.value(Nat.mul(head, Shape.numel(tail)), Shape.Fin.pack(head, Shape.numel(tail), coordinate))) : List<&2, Nat>}:
94 (+first, +last) = coordinate
95 +stride = Shape.numel(tail)
96 +major = Shape.Fin.value(head, first)
97 +minor = Shape.Fin.value(stride, last)
98 %Equal.sym(Nat, Shape.Fin.value(Nat.mul(head, stride), Shape.Fin.pack(head, stride, (first, last))), Nat.add(Nat.mul(major, stride), minor), Index.pack_value(head, stride, first, last)) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == Shape.digits(head <> tail, _) : List<&2, Nat>}
99 %Nats.divmod_exact(major, stride, minor, Shape.Fin.bound(stride, last)) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == Pair.fst(Nat, Nat, _) <> Shape.digits(tail, Pair.snd(Nat, Nat, _)) : List<&2, Nat>}
100 %rec(last) : {major <> Shape.indices(tail, Shape.unflatten(tail, last)) == major <> _ : List<&2, Nat>}
101 {==}
102
103def digits_unflatten(+dims: List<&2, Nat>, +index: Shape.Fin(Shape.numel(dims))) -> {Shape.indices(dims, Shape.unflatten(dims, index)) == Shape.digits(dims, Shape.Fin.value(Shape.numel(dims), index)) : List<&2, Nat>}:
104 match dims:
105 case Nil{}:
106 match index:
107 case Shape.FZero{}:
108 {==}
109 case Shape.FNext{previous}:
110 match previous:
111 case +head <> +tail:
112 %Index.pack_unpack(head, Shape.numel(tail), index) : {Shape.indices(head <> tail, Shape.unflatten(head <> tail, index)) == Shape.digits(head <> tail, Shape.Fin.value(Shape.numel(head <> tail), _)) : List<&2, Nat>}
113 digits_step(head, tail, +index => digits_unflatten(tail, index), Shape.Fin.unpack(head, Shape.numel(tail), index))
114
115def fits_of.step(-dh: Nat, -dt: List<&2, Nat>, -vh: Nat, -vt: List<&2, Nat>, verdict: Shape.BoolCase<Nat.is_lt(vh, dh)>, ok: {True{} == Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)) : Bool}, rec: {True{} == Shape.fit(dt, vt) : Bool} -> Fits(dt, vt)) -> Fits(dh <> dt, vh <> vt):
116 match verdict:
117 case Shape.BFalse{equal}:
118 Empty.absurd(Fits(dh <> dt, vh <> vt), Shape.clash(Equal.trans(Bool, True{}, Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)), False{}, ok, Equal.cong(Bool, Bool, head => Bool.and(head, Shape.fit(dt, vt)), Nat.is_lt(vh, dh), False{}, Equal.sym(Bool, False{}, Nat.is_lt(vh, dh), equal)))))
119 case Shape.BTrue{+equal}:
120 FStep{equal, rec(Equal.trans(Bool, True{}, Bool.and(Nat.is_lt(vh, dh), Shape.fit(dt, vt)), Shape.fit(dt, vt), ok, Equal.cong(Bool, Bool, head => Bool.and(head, Shape.fit(dt, vt)), Nat.is_lt(vh, dh), True{}, Equal.sym(Bool, True{}, Nat.is_lt(vh, dh), equal))))}
121
122def fits_of(+dims: List<&2, Nat>, +values: List<&2, Nat>, ok: {True{} == Shape.fit(dims, values) : Bool}) -> Fits(dims, values):
123 match dims values:
124 case Nil{} Nil{}:
125 Unit{}
126 case +dh <> +dt +vh <> +vt:
127 fits_of.step(dh, dt, vh, vt, Shape.bool_case(Nat.is_lt(vh, dh)), ok, rest => fits_of(dt, vt, rest))
128 case Nil{} vh <> vt:
129 Empty.absurd(Fits(Nil{}, vh <> vt), Shape.clash(ok))
130 case dh <> dt Nil{}:
131 Empty.absurd(Fits(dh <> dt, Nil{}), Shape.clash(ok))
132
133def fit_indices(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {True{} == Shape.fit(dims, Shape.indices(dims, coordinate)) : Bool}:
134 match dims:
135 case Nil{}:
136 match coordinate:
137 case Unit{}:
138 {==}
139 case +head <> +tail:
140 Shape.CStep{+index, +rest} = coordinate
141 %Nats.lt_true(Shape.Fin.value(head, index), head, Shape.Fin.bound(head, index)) : {True{} == Bool.and(_, Shape.fit(tail, Shape.indices(tail, rest))) : Bool}
142 fit_indices(tail, rest)
143
144def fit_words_spec(+dims: List<&2, U32>, +values: List<&2, U32>) -> {Shape.fit_words(dims, values) == Shape.fit(Shape.naturals(dims), Shape.naturals(values)) : Bool}:
145 match dims values:
146 case Nil{} Nil{}:
147 {==}
148 case +dh <> +dt +vh <> +vt:
149 Equal.cong(Bool, Bool, rest => Bool.and(Nat.is_lt(U32.to_nat(vh), U32.to_nat(dh)), rest), Shape.fit_words(dt, vt), Shape.fit(Shape.naturals(dt), Shape.naturals(vt)), fit_words_spec(dt, vt))
150 case Nil{} vh <> vt:
151 {==}
152 case dh <> dt Nil{}:
153 {==}
154
155def digits_words_spec(+dims: List<&2, U32>, +index: Nat) -> {Shape.digits_words(dims, index) == Shape.digits(Shape.naturals(dims), index) : List<&2, Nat>}:
156 match dims:
157 case Nil{}:
158 {==}
159 case +head <> +tail:
160 +stride = Shape.numel(Shape.naturals(tail))
161 %Equal.sym(Nat, Shape.numel_words(tail), stride, Shape.numel_words_spec(tail)) : {Nat.div(index, _) <> Shape.digits_words(tail, Nat.mod(index, _)) == Shape.digits(Shape.naturals(head <> tail), index) : List<&2, Nat>}
162 Equal.cong(List<&2, Nat>, List<&2, Nat>, rest => Nat.div(index, stride) <> rest, Shape.digits_words(tail, Nat.mod(index, stride)), Shape.digits(Shape.naturals(tail), Nat.mod(index, stride)), digits_words_spec(tail, Nat.mod(index, stride)))
163
164def projected_words_spec(+dims: List<&2, U32>, +values: List<&2, U32>) -> {Shape.projected_words(dims, values) == Shape.broadcast_indices(Shape.naturals(dims), Shape.naturals(values)) : List<&2, Nat>}:
165 match dims values:
166 case +head <> +tail +value <> +rest:
167 Equal.cong(List<&2, Nat>, List<&2, Nat>, rest => Shape.broadcast_index(U32.to_nat(head), U32.to_nat(value)) <> rest, Shape.projected_words(tail, rest), Shape.broadcast_indices(Shape.naturals(tail), Shape.naturals(rest)), projected_words_spec(tail, rest))
168 case Nil{} values:
169 {==}
170 case head <> tail Nil{}:
171 {==}
172
173def naturals_drop(+values: List<&2, U32>, +count: Nat) -> {Shape.naturals(List.drop(&2, U32, values, count)) == List.drop(&2, Nat, Shape.naturals(values), count) : List<&2, Nat>}:
174 match values count:
175 case Nil{} count:
176 {==}
177 case head <> tail 0n:
178 {==}
179 case head <> tail 1n+rest:
180 naturals_drop(tail, rest)
181
182def horner_spec(+dims: List<&2, U32>, +values: List<&2, U32>, +acc: Nat, +fits: Fits(Shape.naturals(dims), Shape.naturals(values))) -> {Shape.horner(dims, values, acc) == Nat.add(Nat.mul(acc, Shape.numel(Shape.naturals(dims))), Shape.offset(Shape.naturals(dims), Shape.naturals(values))) : Nat}:
183 match dims values:
184 case Nil{} Nil{}:
185 Equal.trans(Nat, acc, Nat.mul(acc, 1n), Nat.add(Nat.mul(acc, 1n), 0n), Equal.sym(Nat, Nat.mul(acc, 1n), acc, Nats.mul_one_right(acc)), Equal.sym(Nat, Nat.add(Nat.mul(acc, 1n), 0n), Nat.mul(acc, 1n), Broadcast.add_zero(Nat.mul(acc, 1n))))
186 case +dh <> +dt +vh <> +vt:
187 FStep{bound, +tail} = fits
188 +size = Shape.numel(Shape.naturals(dt))
189 +rest = Shape.offset(Shape.naturals(dt), Shape.naturals(vt))
190 +next = Nat.add(Nat.mul(acc, U32.to_nat(dh)), U32.to_nat(vh))
191 %Equal.sym(Nat, Shape.horner(dt, vt, next), Nat.add(Nat.mul(next, size), rest), horner_spec(dt, vt, next, tail)) : {_ == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
192 %Equal.sym(Nat, Nat.mul(next, size), Nat.add(Nat.mul(Nat.mul(acc, U32.to_nat(dh)), size), Nat.mul(U32.to_nat(vh), size)), Nats.mul_add_right(Nat.mul(acc, U32.to_nat(dh)), U32.to_nat(vh), size)) : {Nat.add(_, rest) == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
193 %Equal.sym(Nat, Nat.mul(Nat.mul(acc, U32.to_nat(dh)), size), Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nats.mul_assoc(acc, U32.to_nat(dh), size)) : {Nat.add(Nat.add(_, Nat.mul(U32.to_nat(vh), size)), rest) == Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)) : Nat}
194 Equal.sym(Nat, Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.add(Nat.mul(U32.to_nat(vh), size), rest)), Nat.add(Nat.add(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.mul(U32.to_nat(vh), size)), rest), Index.add_assoc(Nat.mul(acc, Nat.mul(U32.to_nat(dh), size)), Nat.mul(U32.to_nat(vh), size), rest))
195 case Nil{} vh <> vt:
196 match fits:
197 case dh <> dt Nil{}:
198 match fits:
199
200def pair_fits(+shape: Shape.Shape, +first: U32, +second: U32, +x: U32, +y: U32,
201 dimensions: {Shape.dimensions(shape) == [first, second] : List<&2, U32>},
202 first_inside: {Nat.is_lt(U32.to_nat(x), U32.to_nat(first)) == True{} : Bool},
203 second_inside: {Nat.is_lt(U32.to_nat(y), U32.to_nat(second)) == True{} : Bool}) -> Fits(Shape.extents(shape), [U32.to_nat(x), U32.to_nat(y)]):
204 fits_of(Shape.extents(shape), [U32.to_nat(x), U32.to_nat(y)],
205 %Equal.sym(List<&2, U32>, Shape.dimensions(shape), [first, second], dimensions) :
206 {True{} == Shape.fit(Shape.naturals(_), [U32.to_nat(x), U32.to_nat(y)]) : Bool}
207 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(x), U32.to_nat(first)), True{}, first_inside) :
208 {True{} == Bool.and(_, Bool.and(Nat.is_lt(U32.to_nat(y), U32.to_nat(second)), True{})) : Bool}
209 %Equal.sym(Bool, Nat.is_lt(U32.to_nat(y), U32.to_nat(second)), True{}, second_inside) :
210 {True{} == Bool.and(True{}, Bool.and(_, True{})) : Bool}
211 {==})