src/shape/plan_proofs.bend
bendtorch-core@0.1.0.0 · 191 lines · 21 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/broadcast_proofs.bend as Broadcast
5import bendtorch-core@0.1.0.0/src/shape/unique_proofs.bend as Unique
6import bendtorch-core@0.1.0.0/src/shape/construct_proofs.bend as Construct
7import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as Laws
8
9def plan_construct(+left: Shape.Shape, +right: Shape.Shape, -left_padding: Nat, -right_padding: Nat, -minimal: Shape.PadMinimal<left_padding, right_padding>, -output: List<&2, U32>, -aligned: Shape.Aligned(Shape.pad(left_padding, Shape.extents(left)), Shape.pad(right_padding, Shape.extents(right)), Shape.naturals(output)), +result: Shape.ConstructResult<output>) -> {Laws.planned_dimensions(left, right, Shape.broadcast_checked.construct(left, right, left_padding, right_padding, minimal, output, aligned, result)) == (Shape.dimensions(left), Shape.dimensions(right)) : List<&2, U32> & List<&2, U32>}:
10 match result:
11 case Shape.SError{error}:
12 {==}
13 case Shape.SBuilt{shape, preserved}:
14 {==}
15
16def plan_align(+left: Shape.Shape, +right: Shape.Shape, +left_padding: Nat, +right_padding: Nat, -minimal: Shape.PadMinimal<left_padding, right_padding>, +alignment: Shape.Alignment<Shape.pad_words(left_padding, Shape.dimensions(left)), Shape.pad_words(right_padding, Shape.dimensions(right))>) -> {Laws.planned_dimensions(left, right, Shape.broadcast_checked.align(left, right, left_padding, right_padding, minimal, alignment)) == (Shape.dimensions(left), Shape.dimensions(right)) : List<&2, U32> & List<&2, U32>}:
17 match alignment:
18 case Shape.Misaligned{}:
19 {==}
20 case Shape.Aligns{+output, aligned}:
21 plan_construct(left, right, left_padding, right_padding, minimal, output, Shape.broadcast_checked.repad(left_padding, right_padding, Shape.dimensions(left), Shape.dimensions(right), output, aligned), Shape.construct(output))
22
23def broadcast_checked_plan(+left: Shape.Shape, +right: Shape.Shape) -> {Laws.planned_dimensions(left, right, Shape.broadcast_checked(left, right)) == (Shape.dimensions(left), Shape.dimensions(right)) : List<&2, U32> & List<&2, U32>}:
24 +left_rank = Shape.rank_dimensions(Shape.dimensions(left))
25 +right_rank = Shape.rank_dimensions(Shape.dimensions(right))
26 plan_align(left, right, Nat.sub(right_rank, left_rank), Nat.sub(left_rank, right_rank), Shape.minimal(left_rank, right_rank), Shape.align(Shape.pad_words(Nat.sub(right_rank, left_rank), Shape.dimensions(left)), Shape.pad_words(Nat.sub(left_rank, right_rank), Shape.dimensions(right))))
27
28def some_injective(-A: Data, -x: A, -y: A, equal: {Some{x} == Some{y} : Maybe<&2, A>}) -> {x == y : A}:
29 Equal.cong(Maybe<&2, A>, A, found => Maybe.default(&2, A, found, x), Some{x}, Some{y}, equal)
30
31def some_none(-A: Data, -x: A, equal: {Some{x} == None{} : Maybe<&2, A>}) -> Empty:
32 Shape.clash(Equal.cong(Maybe<&2, A>, Bool, found => Maybe.is_some(&2, A, found), Some{x}, None{}, equal))
33
34def is_eq_refl(+value: Nat) -> {True{} == Nat.is_eq(value, value) : Bool}:
35 match value:
36 case 0n:
37 {==}
38 case 1n+p:
39 is_eq_refl(p)
40
41def eq_true(+a: Nat, -b: Nat, equal: {a == b : Nat}) -> {True{} == Nat.is_eq(a, b) : Bool}:
42 %equal : {True{} == Nat.is_eq(a, _) : Bool}
43 is_eq_refl(a)
44
45def sub_balance(+a: Nat, +b: Nat) -> {Nat.add(Nat.sub(b, a), a) == Nat.add(Nat.sub(a, b), b) : Nat}:
46 match a b:
47 case 0n 0n:
48 {==}
49 case 0n 1n+q:
50 Broadcast.add_zero(1n+q)
51 case 1n+p 0n:
52 Equal.sym(Nat, Nat.add(1n+p, 0n), 1n+p, Broadcast.add_zero(1n+p))
53 case 1n+p 1n+q:
54 %Broadcast.add_succ(Nat.sub(q, p), p) : {_ == Nat.add(Nat.sub(p, q), 1n+q) : Nat}
55 %Broadcast.add_succ(Nat.sub(p, q), q) : {1n+Nat.add(Nat.sub(q, p), p) == _ : Nat}
56 Equal.cong(Nat, Nat, value => 1n+value, Nat.add(Nat.sub(q, p), p), Nat.add(Nat.sub(p, q), q), sub_balance(p, q))
57
58def choice_output(-left: U32, -right: U32, choice: Shape.AxisChoice<left, right>) -> Maybe<&2, Nat>:
59 match choice:
60 case Shape.Picks{output, axis}:
61 Some{U32.to_nat(output)}
62 case Shape.Clashes{}:
63 None{}
64
65def aligned_output(-left: List<&2, U32>, -right: List<&2, U32>, alignment: Shape.Alignment<left, right>) -> Maybe<&2, List<&2, Nat>>:
66 match alignment:
67 case Shape.Aligns{output, aligned}:
68 Some{Shape.naturals(output)}
69 case Shape.Misaligned{}:
70 None{}
71
72def some_nat(-a: Nat, -b: Nat, equal: {a == b : Nat}) -> {Some{a} == Some{b} : Maybe<&2, Nat>}:
73 Equal.cong(Nat, Maybe<&2, Nat>, value => Some{value}, a, b, equal)
74
75def axis_right_complete(+left: U32, +right: U32, -output: Nat, +one: Shape.BoolCase<Nat.is_eq(1n, U32.to_nat(right))>, +axis: Shape.Axis<U32.to_nat(left), U32.to_nat(right), output>, differ: {False{} == Nat.is_eq(U32.to_nat(left), U32.to_nat(right)) : Bool}, left_wide: {False{} == Nat.is_eq(1n, U32.to_nat(left)) : Bool}) -> {Some{output} == choice_output(left, right, Shape.axis_choice.right_one(left, right, one)) : Maybe<&2, Nat>}:
76 match one:
77 case Shape.BTrue{right_one}:
78 match axis:
79 case Shape.AxisEqual{left_eq, right_eq}:
80 some_nat(output, U32.to_nat(left), left_eq)
81 case Shape.AxisLeftOne{left_one, right_eq}:
82 Empty.absurd({Some{output} == Some{U32.to_nat(left)} : Maybe<&2, Nat>}, Nats.contradiction(Nat.is_eq(1n, U32.to_nat(left)), eq_true(1n, U32.to_nat(left), left_one), left_wide))
83 case Shape.AxisRightOne{left_eq, right_one_eq}:
84 some_nat(output, U32.to_nat(left), left_eq)
85 case Shape.BFalse{right_wide}:
86 match axis:
87 case Shape.AxisEqual{left_eq, right_eq}:
88 Empty.absurd({Some{output} == None{} : Maybe<&2, Nat>}, Nats.contradiction(Nat.is_eq(U32.to_nat(left), U32.to_nat(right)), eq_true(U32.to_nat(left), U32.to_nat(right), Equal.trans(Nat, U32.to_nat(left), output, U32.to_nat(right), Equal.sym(Nat, output, U32.to_nat(left), left_eq), right_eq)), differ))
89 case Shape.AxisLeftOne{left_one, right_eq}:
90 Empty.absurd({Some{output} == None{} : Maybe<&2, Nat>}, Nats.contradiction(Nat.is_eq(1n, U32.to_nat(left)), eq_true(1n, U32.to_nat(left), left_one), left_wide))
91 case Shape.AxisRightOne{left_eq, right_one}:
92 Empty.absurd({Some{output} == None{} : Maybe<&2, Nat>}, Nats.contradiction(Nat.is_eq(1n, U32.to_nat(right)), eq_true(1n, U32.to_nat(right), right_one), right_wide))
93
94def axis_left_complete(+left: U32, +right: U32, -output: Nat, +one: Shape.BoolCase<Nat.is_eq(1n, U32.to_nat(left))>, +axis: Shape.Axis<U32.to_nat(left), U32.to_nat(right), output>, differ: {False{} == Nat.is_eq(U32.to_nat(left), U32.to_nat(right)) : Bool}) -> {Some{output} == choice_output(left, right, Shape.axis_choice.left_one(left, right, one)) : Maybe<&2, Nat>}:
95 match one:
96 case Shape.BTrue{left_one}:
97 match axis:
98 case Shape.AxisEqual{left_eq, right_eq}:
99 some_nat(output, U32.to_nat(right), right_eq)
100 case Shape.AxisLeftOne{left_one_eq, right_eq}:
101 some_nat(output, U32.to_nat(right), right_eq)
102 case Shape.AxisRightOne{left_eq, right_one}:
103 some_nat(output, U32.to_nat(right), Equal.trans(Nat, output, 1n, U32.to_nat(right), Equal.trans(Nat, output, U32.to_nat(left), 1n, left_eq, Equal.sym(Nat, 1n, U32.to_nat(left), Shape.equal_decided(1n, U32.to_nat(left), left_one))), right_one))
104 case Shape.BFalse{left_wide}:
105 axis_right_complete(left, right, output, Shape.bool_case(Nat.is_eq(1n, U32.to_nat(right))), axis, differ, left_wide)
106
107def axis_equal_complete(+left: U32, +right: U32, -output: Nat, +same: Shape.BoolCase<Nat.is_eq(U32.to_nat(left), U32.to_nat(right))>, +axis: Shape.Axis<U32.to_nat(left), U32.to_nat(right), output>) -> {Some{output} == choice_output(left, right, Shape.axis_choice.equal(left, right, same)) : Maybe<&2, Nat>}:
108 match same:
109 case Shape.BTrue{equal}:
110 match axis:
111 case Shape.AxisEqual{left_eq, right_eq}:
112 some_nat(output, U32.to_nat(left), left_eq)
113 case Shape.AxisLeftOne{left_one, right_eq}:
114 some_nat(output, U32.to_nat(left), Equal.trans(Nat, output, U32.to_nat(right), U32.to_nat(left), right_eq, Equal.sym(Nat, U32.to_nat(left), U32.to_nat(right), Shape.equal_decided(U32.to_nat(left), U32.to_nat(right), equal))))
115 case Shape.AxisRightOne{left_eq, right_one}:
116 some_nat(output, U32.to_nat(left), left_eq)
117 case Shape.BFalse{differ}:
118 axis_left_complete(left, right, output, Shape.bool_case(Nat.is_eq(1n, U32.to_nat(left))), axis, differ)
119
120def axis_complete(+left: U32, +right: U32, -output: Nat, +axis: Shape.Axis<U32.to_nat(left), U32.to_nat(right), output>) -> {Some{output} == choice_output(left, right, Shape.axis_choice(left, right)) : Maybe<&2, Nat>}:
121 axis_equal_complete(left, right, output, Shape.bool_case(Nat.is_eq(U32.to_nat(left), U32.to_nat(right))), axis)
122
123def align_step_complete(-lh: U32, -lt: List<&2, U32>, -rh: U32, -rt: List<&2, U32>, -oh: Nat, -ot: List<&2, Nat>, +rest: Shape.Alignment<lt, rt>, +choice: Shape.AxisChoice<lh, rh>, rest_seen: {Some{ot} == aligned_output(lt, rt, rest) : Maybe<&2, List<&2, Nat>>}, choice_seen: {Some{oh} == choice_output(lh, rh, choice) : Maybe<&2, Nat>}) -> {Some{oh <> ot} == aligned_output(lh <> lt, rh <> rt, Shape.align.step(lh, lt, rh, rt, rest, choice)) : Maybe<&2, List<&2, Nat>>}:
124 match rest:
125 case Shape.Misaligned{}:
126 Empty.absurd({Some{oh <> ot} == aligned_output(lh <> lt, rh <> rt, Shape.align.step(lh, lt, rh, rt, Shape.Misaligned{}, choice)) : Maybe<&2, List<&2, Nat>>}, some_none(List<&2, Nat>, ot, rest_seen))
127 case Shape.Aligns{tail, aligned}:
128 match choice:
129 case Shape.Clashes{}:
130 Empty.absurd({Some{oh <> ot} == None{} : Maybe<&2, List<&2, Nat>>}, some_none(Nat, oh, choice_seen))
131 case Shape.Picks{output, axis}:
132 %some_injective(Nat, oh, U32.to_nat(output), choice_seen) : {Some{oh <> ot} == Some{_ <> Shape.naturals(tail)} : Maybe<&2, List<&2, Nat>>}
133 %some_injective(List<&2, Nat>, ot, Shape.naturals(tail), rest_seen) : {Some{oh <> ot} == Some{oh <> _} : Maybe<&2, List<&2, Nat>>}
134 {==}
135
136def align_complete(+left: List<&2, U32>, +right: List<&2, U32>, +output: List<&2, Nat>, +aligned: Shape.Aligned(Shape.naturals(left), Shape.naturals(right), output)) -> {Some{output} == aligned_output(left, right, Shape.align(left, right)) : Maybe<&2, List<&2, Nat>>}:
137 match left right output:
138 case Nil{} Nil{} Nil{}:
139 {==}
140 case +lh <> +lt +rh <> +rt +oh <> +ot:
141 Shape.AStep{+axis, +tail} = aligned
142 align_step_complete(lh, lt, rh, rt, oh, ot, Shape.align(lt, rt), Shape.axis_choice(lh, rh), align_complete(lt, rt, ot, tail), axis_complete(lh, rh, oh, axis))
143 case Nil{} Nil{} oh <> ot:
144 match aligned:
145 case Nil{} rh <> rt Nil{}:
146 match aligned:
147 case lh <> lt Nil{} Nil{}:
148 match aligned:
149 case Nil{} rh <> rt oh <> ot:
150 match aligned:
151 case lh <> lt Nil{} oh <> ot:
152 match aligned:
153 case lh <> lt rh <> rt Nil{}:
154 match aligned:
155
156def complete_construct(+left: Shape.Shape, +right: Shape.Shape, -left_padding: Nat, -right_padding: Nat, -minimal: Shape.PadMinimal<left_padding, right_padding>, -output: List<&2, U32>, -aligned: Shape.Aligned(Shape.pad(left_padding, Shape.extents(left)), Shape.pad(right_padding, Shape.extents(right)), Shape.naturals(output)), +result: Shape.ConstructResult<output>, built: {Laws.built(output, result) == Some{output} : Maybe<&2, List<&2, U32>>}) -> {Laws.broadcast_built(Shape.broadcast_checked.construct(left, right, left_padding, right_padding, minimal, output, aligned, result)) == True{} : Bool}:
157 match result:
158 case Shape.SError{error}:
159 Empty.absurd({False{} == True{} : Bool}, some_none(List<&2, U32>, output, Equal.sym(Maybe<&2, List<&2, U32>>, None{}, Some{output}, built)))
160 case Shape.SBuilt{shape, preserved}:
161 {==}
162
163def complete_align(+left: Shape.Shape, +right: Shape.Shape, +left_padding: Nat, +right_padding: Nat, -minimal: Shape.PadMinimal<left_padding, right_padding>, +alignment: Shape.Alignment<Shape.pad_words(left_padding, Shape.dimensions(left)), Shape.pad_words(right_padding, Shape.dimensions(right))>, -target: List<&2, Nat>, seen: {Some{target} == aligned_output(Shape.pad_words(left_padding, Shape.dimensions(left)), Shape.pad_words(right_padding, Shape.dimensions(right)), alignment) : Maybe<&2, List<&2, Nat>>}, +rank_bound: Shape.LE(List.length(&2, Nat, target), 4n), +capacity_bound: Shape.LE(Shape.numel(target), 1048576n)) -> {Laws.broadcast_built(Shape.broadcast_checked.align(left, right, left_padding, right_padding, minimal, alignment)) == True{} : Bool}:
164 match alignment:
165 case Shape.Misaligned{}:
166 Empty.absurd({False{} == True{} : Bool}, some_none(List<&2, Nat>, target, seen))
167 case Shape.Aligns{+output, aligned}:
168 +same = some_injective(List<&2, Nat>, target, Shape.naturals(output), seen)
169 +rank_fits = Nats.le_rewrite(Shape.rank_dimensions(output), List.length(&2, Nat, target), 4n, Equal.trans(Nat, List.length(&2, Nat, target), List.length(&2, Nat, Shape.naturals(output)), Shape.rank_dimensions(output), Equal.cong(List<&2, Nat>, Nat, values => List.length(&2, Nat, values), target, Shape.naturals(output), same), Equal.sym(Nat, Shape.rank_dimensions(output), List.length(&2, Nat, Shape.naturals(output)), Shape.rank_naturals(output))), rank_bound)
170 +capacity_fits = Nats.le_rewrite(Shape.numel(Shape.naturals(output)), Shape.numel(target), 1048576n, Equal.cong(List<&2, Nat>, Nat, values => Shape.numel(values), target, Shape.naturals(output), same), capacity_bound)
171 complete_construct(left, right, left_padding, right_padding, minimal, output, Shape.broadcast_checked.repad(left_padding, right_padding, Shape.dimensions(left), Shape.dimensions(right), output, aligned), Shape.construct(output), Construct.construction_complete(output, Shape.bool_case(Nat.is_le(Shape.rank_dimensions(output), 4n)), Nats.le_true(Shape.rank_dimensions(output), 4n, rank_fits), capacity_fits))
172
173def padded_words(+left_padding: Nat, +right_padding: Nat, -left: List<&2, U32>, -right: List<&2, U32>, -output: List<&2, Nat>, aligned: Shape.Aligned(Shape.pad(left_padding, Shape.naturals(left)), Shape.pad(right_padding, Shape.naturals(right)), output)) -> Shape.Aligned(Shape.naturals(Shape.pad_words(left_padding, left)), Shape.naturals(Shape.pad_words(right_padding, right)), output):
174 %Equal.sym(List<&2, Nat>, Shape.naturals(Shape.pad_words(left_padding, left)), Shape.pad(left_padding, Shape.naturals(left)), Shape.naturals_pad(left_padding, left)) : Shape.Aligned(_, Shape.naturals(Shape.pad_words(right_padding, right)), output)
175 %Equal.sym(List<&2, Nat>, Shape.naturals(Shape.pad_words(right_padding, right)), Shape.pad(right_padding, Shape.naturals(right)), Shape.naturals_pad(right_padding, right)) : Shape.Aligned(Shape.pad(left_padding, Shape.naturals(left)), _, output)
176 aligned
177
178def broadcast_checked_complete(+left: Shape.Shape, +right: Shape.Shape, +output: List<&2, U32>, +alignment: Shape.Broadcast<Shape.extents(left), Shape.extents(right), Shape.naturals(output)>, +rank_bound: Shape.LE(Shape.rank_dimensions(output), 4n), +capacity_bound: Shape.LE(Shape.logical_numel(output), 1048576n)) -> {Laws.broadcast_built(Shape.broadcast_checked(left, right)) == True{} : Bool}:
179 match alignment:
180 case Shape.BAlign{+left_padding, +right_padding, +minimal, +aligned}:
181 +left_rank = Shape.rank_dimensions(Shape.dimensions(left))
182 +right_rank = Shape.rank_dimensions(Shape.dimensions(right))
183 +left_extra = Nat.sub(right_rank, left_rank)
184 +right_extra = Nat.sub(left_rank, right_rank)
185 +balanced = Equal.trans(Nat, Nat.add(left_padding, left_rank), Nat.add(left_padding, List.length(&2, Nat, Shape.extents(left))), Nat.add(right_padding, right_rank), Equal.cong(Nat, Nat, value => Nat.add(left_padding, value), left_rank, List.length(&2, Nat, Shape.extents(left)), Shape.rank_naturals(Shape.dimensions(left))), Equal.trans(Nat, Nat.add(left_padding, List.length(&2, Nat, Shape.extents(left))), Nat.add(right_padding, List.length(&2, Nat, Shape.extents(right))), Nat.add(right_padding, right_rank), Unique.balance(left_padding, right_padding, Shape.extents(left), Shape.extents(right), Shape.naturals(output), aligned), Equal.sym(Nat, Nat.add(right_padding, right_rank), Nat.add(right_padding, List.length(&2, Nat, Shape.extents(right))), Equal.cong(Nat, Nat, value => Nat.add(right_padding, value), right_rank, List.length(&2, Nat, Shape.extents(right)), Shape.rank_naturals(Shape.dimensions(right))))))
186 +computed = sub_balance(left_rank, right_rank)
187 +same_left = Unique.left_padding_equal(left_padding, right_padding, left_extra, right_extra, left_rank, right_rank, minimal, Shape.minimal(left_rank, right_rank), balanced, computed)
188 +same_right = Unique.cancel_right(right_padding, right_extra, right_rank, Equal.trans(Nat, Nat.add(right_padding, right_rank), Nat.add(left_padding, left_rank), Nat.add(right_extra, right_rank), Equal.sym(Nat, Nat.add(left_padding, left_rank), Nat.add(right_padding, right_rank), balanced), Equal.trans(Nat, Nat.add(left_padding, left_rank), Nat.add(left_extra, left_rank), Nat.add(right_extra, right_rank), Equal.cong(Nat, Nat, value => Nat.add(value, left_rank), left_padding, left_extra, same_left), computed)))
189 +retyped = Unique.retype(left_extra, left_padding, right_extra, right_padding, Shape.extents(left), Shape.extents(right), Shape.naturals(output), Equal.sym(Nat, left_padding, left_extra, same_left), Equal.sym(Nat, right_padding, right_extra, same_right), aligned)
190 +seen = align_complete(Shape.pad_words(left_extra, Shape.dimensions(left)), Shape.pad_words(right_extra, Shape.dimensions(right)), Shape.naturals(output), padded_words(left_extra, right_extra, Shape.dimensions(left), Shape.dimensions(right), Shape.naturals(output), retyped))
191 complete_align(left, right, left_extra, right_extra, Shape.minimal(left_rank, right_rank), Shape.align(Shape.pad_words(left_extra, Shape.dimensions(left)), Shape.pad_words(right_extra, Shape.dimensions(right))), Shape.naturals(output), seen, Nats.le_rewrite(List.length(&2, Nat, Shape.naturals(output)), Shape.rank_dimensions(output), 4n, Shape.rank_naturals(output), rank_bound), capacity_bound)