src/shape/index_proofs.bend
bendtorch-core@0.1.0.0 · 154 lines · 16 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3
4def add_assoc(a: Nat, -b: Nat, -c: Nat) -> {Nat.add(a, Nat.add(b, c)) == Nat.add(Nat.add(a, b), c) : Nat}:
5 match a:
6 case 0n:
7 {==}
8 case 1n+p:
9 %add_assoc(p, b, c) : {1n+Nat.add(p, Nat.add(b, c)) == 1n+_ : Nat}
10 {==}
11
12def split_left(+a: Nat, +b: Nat, +index: Shape.Fin(a)) -> {Shape.Fin.split(a, b, Shape.Fin.left(a, b, index)) == Inl{index} : Either<Shape.Fin(a), Shape.Fin(b)>}:
13 match a:
14 case 0n:
15 match index:
16 case 1n+p:
17 match index:
18 case Shape.FZero{}:
19 {==}
20 case Shape.FNext{+previous}:
21 %split_left(p, b, previous) : {Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, Shape.Fin.left(p, b, previous))) == Shape.Fin.split.step(p, b, _) : Either<Shape.Fin(1n+p), Shape.Fin(b)>}
22 {==}
23
24def split_shift(+a: Nat, +b: Nat, +index: Shape.Fin(b)) -> {Shape.Fin.split(a, b, Shape.Fin.shift(a, b, index)) == Inr{index} : Either<Shape.Fin(a), Shape.Fin(b)>}:
25 match a:
26 case 0n:
27 {==}
28 case 1n+p:
29 %split_shift(p, b, index) : {Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, Shape.Fin.shift(p, b, index))) == Shape.Fin.split.step(p, b, _) : Either<Shape.Fin(1n+p), Shape.Fin(b)>}
30 {==}
31
32def merge_step(+a: Nat, +b: Nat, split: Either<Shape.Fin(a), Shape.Fin(b)>) -> {Shape.Fin.merge(1n+a, b, Shape.Fin.split.step(a, b, split)) == Shape.FNext{Shape.Fin.merge(a, b, split)} : Shape.Fin(1n+Nat.add(a, b))}:
33 match split:
34 case Inl{index}:
35 {==}
36 case Inr{index}:
37 {==}
38
39def merge_split(+a: Nat, +b: Nat, +index: Shape.Fin(Nat.add(a, b))) -> {Shape.Fin.merge(a, b, Shape.Fin.split(a, b, index)) == index : Shape.Fin(Nat.add(a, b))}:
40 match a:
41 case 0n:
42 {==}
43 case 1n+p:
44 match index:
45 case Shape.FZero{}:
46 {==}
47 case Shape.FNext{+previous}:
48 %merge_split(p, b, previous) : {Shape.Fin.merge(1n+p, b, Shape.Fin.split.step(p, b, Shape.Fin.split(p, b, previous))) == Shape.FNext{_} : Shape.Fin(1n+Nat.add(p, b))}
49 merge_step(p, b, Shape.Fin.split(p, b, previous))
50
51def left_value(+a: Nat, +b: Nat, +index: Shape.Fin(a)) -> {Shape.Fin.value(Nat.add(a, b), Shape.Fin.left(a, b, index)) == Shape.Fin.value(a, index) : Nat}:
52 match a:
53 case 0n:
54 match index:
55 case 1n+p:
56 match index:
57 case Shape.FZero{}:
58 {==}
59 case Shape.FNext{+previous}:
60 %left_value(p, b, previous) : {1n+Shape.Fin.value(Nat.add(p, b), Shape.Fin.left(p, b, previous)) == 1n+_ : Nat}
61 {==}
62
63def shift_value(+a: Nat, +b: Nat, +index: Shape.Fin(b)) -> {Shape.Fin.value(Nat.add(a, b), Shape.Fin.shift(a, b, index)) == Nat.add(a, Shape.Fin.value(b, index)) : Nat}:
64 match a:
65 case 0n:
66 {==}
67 case 1n+p:
68 %shift_value(p, b, index) : {1n+Shape.Fin.value(Nat.add(p, b), Shape.Fin.shift(p, b, index)) == 1n+_ : Nat}
69 {==}
70
71def pack_value(+a: Nat, +b: Nat, +first: Shape.Fin(a), +last: Shape.Fin(b)) -> {Shape.Fin.value(Nat.mul(a, b), Shape.Fin.pack(a, b, (first, last))) == Nat.add(Nat.mul(Shape.Fin.value(a, first), b), Shape.Fin.value(b, last)) : Nat}:
72 match a:
73 case 0n:
74 match first:
75 case 1n+p:
76 match first:
77 case Shape.FZero{}:
78 left_value(b, Nat.mul(p, b), last)
79 case Shape.FNext{+previous}:
80 %add_assoc(b, Nat.mul(Shape.Fin.value(p, previous), b), Shape.Fin.value(b, last)) : {Shape.Fin.value(Nat.add(b, Nat.mul(p, b)), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))) == _ : Nat}
81 %pack_value(p, b, previous, last) : {Shape.Fin.value(Nat.add(b, Nat.mul(p, b)), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))) == Nat.add(b, _) : Nat}
82 shift_value(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last)))
83
84def unpack_pack(+a: Nat, +b: Nat, +first: Shape.Fin(a), +last: Shape.Fin(b)) -> {Shape.Fin.unpack(a, b, Shape.Fin.pack(a, b, (first, last))) == (first, last) : Shape.Fin(a) & Shape.Fin(b)}:
85 match a:
86 case 0n:
87 match first:
88 case 1n+p:
89 match first:
90 case Shape.FZero{}:
91 %split_left(b, Nat.mul(p, b), last) : {Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), Shape.Fin.left(b, Nat.mul(p, b), last))) == Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), _) : Shape.Fin(1n+p) & Shape.Fin(b)}
92 {==}
93 case Shape.FNext{+previous}:
94 %unpack_pack(p, b, previous, last) : {Shape.Fin.unpack(1n+p, b, Shape.Fin.pack(1n+p, b, (Shape.FNext{previous}, last))) == Shape.Fin.unpack.lift(p, b, _) : Shape.Fin(1n+p) & Shape.Fin(b)}
95 %split_shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last))) : {Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), Shape.Fin.shift(b, Nat.mul(p, b), Shape.Fin.pack(p, b, (previous, last))))) == Shape.Fin.unpack.step(p, b, index => Shape.Fin.unpack(p, b, index), _) : Shape.Fin(1n+p) & Shape.Fin(b)}
96 {==}
97
98def pack_lift(+a: Nat, +b: Nat, coordinate: Shape.Fin(a) & Shape.Fin(b)) -> {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.lift(a, b, coordinate)) == Shape.Fin.shift(b, Nat.mul(a, b), Shape.Fin.pack(a, b, coordinate)) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}:
99 (first, last) = coordinate
100 {==}
101
102def pack_unpack_step(+a: Nat, +b: Nat, rec: @+index: Shape.Fin(Nat.mul(a, b)) -> {Shape.Fin.pack(a, b, Shape.Fin.unpack(a, b, index)) == index : Shape.Fin(Nat.mul(a, b))}, split: Either<Shape.Fin(b), Shape.Fin(Nat.mul(a, b))>) -> {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.step(a, b, index => Shape.Fin.unpack(a, b, index), split)) == Shape.Fin.merge(b, Nat.mul(a, b), split) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}:
103 match split:
104 case Inl{last}:
105 {==}
106 case Inr{+index}:
107 %rec(index) : {Shape.Fin.pack(1n+a, b, Shape.Fin.unpack.lift(a, b, Shape.Fin.unpack(a, b, index))) == Shape.Fin.shift(b, Nat.mul(a, b), _) : Shape.Fin(Nat.add(b, Nat.mul(a, b)))}
108 pack_lift(a, b, Shape.Fin.unpack(a, b, index))
109
110def pack_unpack(+a: Nat, +b: Nat, +index: Shape.Fin(Nat.mul(a, b))) -> {Shape.Fin.pack(a, b, Shape.Fin.unpack(a, b, index)) == index : Shape.Fin(Nat.mul(a, b))}:
111 match a:
112 case 0n:
113 match index:
114 case 1n+p:
115 %merge_split(b, Nat.mul(p, b), index) : {Shape.Fin.pack(1n+p, b, Shape.Fin.unpack(1n+p, b, index)) == _ : Shape.Fin(Nat.add(b, Nat.mul(p, b)))}
116 pack_unpack_step(p, b, +index => pack_unpack(p, b, index), Shape.Fin.split(b, Nat.mul(p, b), index))
117
118def flatten_unflatten_step(+head: Nat, +tail: List<&2, Nat>, rec: @+index: Shape.Fin(Shape.numel(tail)) -> {Shape.flatten(tail, Shape.unflatten(tail, index)) == index : Shape.Fin(Shape.numel(tail))}, coordinate: Shape.Fin(head) & Shape.Fin(Shape.numel(tail))) -> {Shape.flatten(head <> tail, Shape.unflatten.step(head, tail, rest => Shape.unflatten(tail, rest), coordinate)) == Shape.Fin.pack(head, Shape.numel(tail), coordinate) : Shape.Fin(Shape.numel(head <> tail))}:
119 (first, +last) = coordinate
120 %rec(last) : {Shape.Fin.pack(head, Shape.numel(tail), (first, Shape.flatten(tail, Shape.unflatten(tail, last)))) == Shape.Fin.pack(head, Shape.numel(tail), (first, _)) : Shape.Fin(Shape.numel(head <> tail))}
121 {==}
122
123def flatten_unflatten(+dims: List<&2, Nat>, +index: Shape.Fin(Shape.numel(dims))) -> {Shape.flatten(dims, Shape.unflatten(dims, index)) == index : Shape.Fin(Shape.numel(dims))}:
124 match dims:
125 case Nil{}:
126 match index:
127 case Shape.FZero{}:
128 {==}
129 case Shape.FNext{previous}:
130 match previous:
131 case +head <> +tail:
132 %pack_unpack(head, Shape.numel(tail), index) : {Shape.flatten(head <> tail, Shape.unflatten(head <> tail, index)) == _ : Shape.Fin(Shape.numel(head <> tail))}
133 flatten_unflatten_step(head, tail, +last => flatten_unflatten(tail, last), Shape.Fin.unpack(head, Shape.numel(tail), index))
134
135def unflatten_flatten(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.unflatten(dims, Shape.flatten(dims, coordinate)) == coordinate : Shape.Coordinate(dims)}:
136 match dims:
137 case Nil{}:
138 match coordinate:
139 case Unit{}:
140 {==}
141 case +head <> +tail:
142 Shape.CStep{+first, +rest} = coordinate
143 %unflatten_flatten(tail, rest) : {Shape.unflatten(head <> tail, Shape.flatten(head <> tail, Shape.CStep{first, rest})) == Shape.CStep{first, _} : Shape.Coordinate(head <> tail)}
144 %unpack_pack(head, Shape.numel(tail), first, Shape.flatten(tail, rest)) : {Shape.unflatten.step(head, tail, last => Shape.unflatten(tail, last), Shape.Fin.unpack(head, Shape.numel(tail), Shape.Fin.pack(head, Shape.numel(tail), (first, Shape.flatten(tail, rest))))) == Shape.unflatten.step(head, tail, last => Shape.unflatten(tail, last), _) : Shape.Coordinate(head <> tail)}
145 {==}
146
147def flatten_row_major(+dims: List<&2, Nat>, +coordinate: Shape.Coordinate(dims)) -> {Shape.Fin.value(Shape.numel(dims), Shape.flatten(dims, coordinate)) == Shape.row_major(dims, coordinate) : Nat}:
148 match dims:
149 case Nil{}:
150 {==}
151 case +head <> +tail:
152 Shape.CStep{+first, +rest} = coordinate
153 %flatten_row_major(tail, rest) : {Shape.Fin.value(Shape.numel(head <> tail), Shape.flatten(head <> tail, Shape.CStep{first, rest})) == Nat.add(Nat.mul(Shape.Fin.value(head, first), Shape.numel(tail)), _) : Nat}
154 pack_value(head, Shape.numel(tail), first, Shape.flatten(tail, rest))