src/shape/projection_proofs.bend
bendtorch-core@0.1.0.0 · 166 lines · 16 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/coordinate_proofs.bend as Coordinates
7
8def strides_of(+dims: List<&2, Nat>) -> List<&2, Nat>:
9 match dims:
10 case Nil{}:
11 Nil{}
12 case head <> tail:
13 Shape.broadcast_index(head, Shape.numel(tail)) <> strides_of(tail)
14
15def zeros(count: Nat, values: List<&2, Nat>) -> List<&2, Nat>:
16 match count:
17 case 0n:
18 values
19 case 1n+rest:
20 0n <> zeros(rest, values)
21
22def dot(steps: List<&2, Nat>, values: List<&2, Nat>) -> Nat:
23 match steps values:
24 case step <> rest +value <> tail:
25 Nat.add(Nat.mul(step, value), dot(rest, tail))
26 case steps values:
27 0n
28
29def lt_zero_false(+value: Nat) -> {Nat.is_lt(value, 0n) == False{} : Bool}:
30 match value:
31 case 0n:
32 {==}
33 case 1n+p:
34 {==}
35
36def walk_left_rows(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +index: Nat, same: {List.length(&2, Nat, right) == List.length(&2, Nat, output) : Nat}) -> {Shape.walk_left(Shape.index_rows(left, right, output), index) == dot(strides_of(left), Shape.digits(output, index)) : Nat}:
37 match left right output:
38 case +lh <> +lt +rh <> +rt +oh <> +ot:
39 +stride = Shape.numel(ot)
40 Equal.cong(Nat, Nat, rest => Nat.add(Nat.mul(Shape.broadcast_index(lh, Shape.numel(lt)), Nat.div(index, stride)), rest), Shape.walk_left(Shape.index_rows(lt, rt, ot), Nat.mod(index, stride)), dot(strides_of(lt), Shape.digits(ot, Nat.mod(index, stride))), walk_left_rows(lt, rt, ot, Nat.mod(index, stride), Unique.succ_injective(List.length(&2, Nat, rt), List.length(&2, Nat, ot), same)))
41 case Nil{} right output:
42 {==}
43 case lh <> lt Nil{} Nil{}:
44 {==}
45 case lh <> lt Nil{} oh <> ot:
46 Empty.absurd({Shape.walk_left(Shape.index_rows(lh <> lt, Nil{}, oh <> ot), index) == dot(strides_of(lh <> lt), Shape.digits(oh <> ot, index)) : Nat}, Unique.succ_zero(List.length(&2, Nat, ot), Equal.sym(Nat, 0n, 1n+List.length(&2, Nat, ot), same)))
47 case lh <> lt rh <> rt Nil{}:
48 {==}
49
50def walk_right_rows(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +index: Nat, same: {List.length(&2, Nat, left) == List.length(&2, Nat, output) : Nat}) -> {Shape.walk_right(Shape.index_rows(left, right, output), index) == dot(strides_of(right), Shape.digits(output, index)) : Nat}:
51 match left right output:
52 case +lh <> +lt +rh <> +rt +oh <> +ot:
53 +stride = Shape.numel(ot)
54 Equal.cong(Nat, Nat, rest => Nat.add(Nat.mul(Shape.broadcast_index(rh, Shape.numel(rt)), Nat.div(index, stride)), rest), Shape.walk_right(Shape.index_rows(lt, rt, ot), Nat.mod(index, stride)), dot(strides_of(rt), Shape.digits(ot, Nat.mod(index, stride))), walk_right_rows(lt, rt, ot, Nat.mod(index, stride), Unique.succ_injective(List.length(&2, Nat, lt), List.length(&2, Nat, ot), same)))
55 case Nil{} Nil{} Nil{}:
56 {==}
57 case Nil{} rh <> rt Nil{}:
58 {==}
59 case Nil{} right oh <> ot:
60 Empty.absurd({Shape.walk_right(Shape.index_rows(Nil{}, right, oh <> ot), index) == dot(strides_of(right), Shape.digits(oh <> ot, index)) : Nat}, Unique.succ_zero(List.length(&2, Nat, ot), Equal.sym(Nat, 0n, 1n+List.length(&2, Nat, ot), same)))
61 case lh <> lt Nil{} output:
62 {==}
63 case lh <> lt rh <> rt Nil{}:
64 {==}
65
66def pad_strides(+count: Nat, +dims: List<&2, Nat>) -> {strides_of(Shape.pad(count, dims)) == zeros(count, strides_of(dims)) : List<&2, Nat>}:
67 match count:
68 case 0n:
69 {==}
70 case 1n+rest:
71 Equal.cong(List<&2, Nat>, List<&2, Nat>, tail => 0n <> tail, strides_of(Shape.pad(rest, dims)), zeros(rest, strides_of(dims)), pad_strides(rest, dims))
72
73def dot_zeros(+count: Nat, +values: List<&2, Nat>, +steps: List<&2, Nat>) -> {dot(zeros(count, steps), values) == dot(steps, List.drop(&2, Nat, values, count)) : Nat}:
74 match count values:
75 case 0n values:
76 %Broadcast.drop_zero(values) : {dot(steps, values) == dot(steps, _) : Nat}
77 {==}
78 case 1n+rest Nil{}:
79 match steps:
80 case Nil{}:
81 {==}
82 case head <> tail:
83 {==}
84 case 1n+rest +value <> +tail:
85 dot_zeros(rest, tail, steps)
86
87def digits_length(+dims: List<&2, Nat>, +index: Nat) -> {List.length(&2, Nat, Shape.digits(dims, index)) == List.length(&2, Nat, dims) : Nat}:
88 match dims:
89 case Nil{}:
90 {==}
91 case head <> +tail:
92 Equal.cong(Nat, Nat, value => 1n+value, List.length(&2, Nat, Shape.digits(tail, Nat.mod(index, Shape.numel(tail)))), List.length(&2, Nat, tail), digits_length(tail, Nat.mod(index, Shape.numel(tail))))
93
94def axis_term(+size: Nat, +stride: Nat, +value: Nat) -> {Nat.mul(Shape.broadcast_index(size, stride), value) == Nat.mul(Shape.broadcast_index(size, value), stride) : Nat}:
95 match size:
96 case 0n:
97 Nats.mul_comm(stride, value)
98 case 1n+p:
99 match p:
100 case 0n:
101 {==}
102 case 1n+q:
103 Nats.mul_comm(stride, value)
104
105def dot_broadcast(+dims: List<&2, Nat>, +values: List<&2, Nat>) -> {dot(strides_of(dims), values) == Shape.offset(dims, Shape.broadcast_indices(dims, values)) : Nat}:
106 match dims values:
107 case +head <> +tail +value <> +rest:
108 %Equal.sym(Nat, dot(strides_of(tail), rest), Shape.offset(tail, Shape.broadcast_indices(tail, rest)), dot_broadcast(tail, rest)) : {Nat.add(Nat.mul(Shape.broadcast_index(head, Shape.numel(tail)), value), _) == Nat.add(Nat.mul(Shape.broadcast_index(head, value), Shape.numel(tail)), Shape.offset(tail, Shape.broadcast_indices(tail, rest))) : Nat}
109 Equal.cong(Nat, Nat, term => Nat.add(term, Shape.offset(tail, Shape.broadcast_indices(tail, rest))), Nat.mul(Shape.broadcast_index(head, Shape.numel(tail)), value), Nat.mul(Shape.broadcast_index(head, value), Shape.numel(tail)), axis_term(head, Shape.numel(tail), value))
110 case Nil{} values:
111 {==}
112 case head <> tail Nil{}:
113 {==}
114
115def padding_of(-width: Nat, +padding: Nat, +rank: Nat, sum: {Nat.add(padding, rank) == width : Nat}) -> {padding == Nat.sub(width, rank) : Nat}:
116 Equal.trans(Nat, padding, Nat.sub(Nat.add(padding, rank), rank), Nat.sub(width, rank), Broadcast.sub_add(padding, rank), Equal.cong(Nat, Nat, value => Nat.sub(value, rank), Nat.add(padding, rank), width, sum))
117
118def left_padded(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>) -> {List.length(&2, Nat, Shape.pad(Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, left)), left)) == List.length(&2, Nat, output) : Nat}:
119 match alignment:
120 case Shape.BAlign{+lp, +rp, minimal, +aligned}:
121 +width = Broadcast.aligned_length_left(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned)
122 %padding_of(List.length(&2, Nat, output), lp, List.length(&2, Nat, left), Equal.trans(Nat, Nat.add(lp, List.length(&2, Nat, left)), List.length(&2, Nat, Shape.pad(lp, left)), List.length(&2, Nat, output), Broadcast.pad_length(lp, left), width)) : {List.length(&2, Nat, Shape.pad(_, left)) == List.length(&2, Nat, output) : Nat}
123 width
124
125def right_padded(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>) -> {List.length(&2, Nat, Shape.pad(Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, right)), right)) == List.length(&2, Nat, output) : Nat}:
126 match alignment:
127 case Shape.BAlign{+lp, +rp, minimal, +aligned}:
128 +width = Broadcast.aligned_length_right(Shape.pad(lp, left), Shape.pad(rp, right), output, aligned)
129 %padding_of(List.length(&2, Nat, output), rp, List.length(&2, Nat, right), Equal.trans(Nat, Nat.add(rp, List.length(&2, Nat, right)), List.length(&2, Nat, Shape.pad(rp, right)), List.length(&2, Nat, output), Broadcast.pad_length(rp, right), width)) : {List.length(&2, Nat, Shape.pad(_, right)) == List.length(&2, Nat, output) : Nat}
130 width
131
132def walk_left_value(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, +index: Shape.Fin(Shape.numel(output))) -> {Shape.walk_left(Shape.index_table(left, right, output), Shape.Fin.value(Shape.numel(output), index)) == Shape.row_major(left, Shape.project_left(left, right, output, alignment, Shape.unflatten(output, index))) : Nat}:
133 +value = Shape.Fin.value(Shape.numel(output), index)
134 +ds = Shape.digits(output, value)
135 +nl = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, left))
136 +nr = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, right))
137 +coordinate = Shape.unflatten(output, index)
138 +projected = Shape.project_left(left, right, output, alignment, coordinate)
139 +strides = strides_of(left)
140 +dropped = List.drop(&2, Nat, ds, nl)
141 %Equal.sym(Nat, Shape.walk_left(Shape.index_rows(Shape.pad(nl, left), Shape.pad(nr, right), output), value), dot(strides_of(Shape.pad(nl, left)), ds), walk_left_rows(Shape.pad(nl, left), Shape.pad(nr, right), output, value, right_padded(left, right, output, alignment))) : {_ == Shape.row_major(left, projected) : Nat}
142 %Equal.cong(List<&2, Nat>, Nat, steps => dot(steps, ds), zeros(nl, strides), strides_of(Shape.pad(nl, left)), Equal.sym(List<&2, Nat>, strides_of(Shape.pad(nl, left)), zeros(nl, strides), pad_strides(nl, left))) : {_ == Shape.row_major(left, projected) : Nat}
143 %Equal.sym(Nat, dot(zeros(nl, strides), ds), dot(strides, dropped), dot_zeros(nl, ds, strides)) : {_ == Shape.row_major(left, projected) : Nat}
144 %Equal.sym(Nat, dot(strides, dropped), Shape.offset(left, Shape.broadcast_indices(left, dropped)), dot_broadcast(left, dropped)) : {_ == Shape.row_major(left, projected) : Nat}
145 %Equal.cong(Nat, Nat, width => Shape.offset(left, Shape.broadcast_indices(left, List.drop(&2, Nat, ds, Nat.sub(width, List.length(&2, Nat, left))))), List.length(&2, Nat, ds), List.length(&2, Nat, output), digits_length(output, value)) : {_ == Shape.row_major(left, projected) : Nat}
146 %Equal.cong(List<&2, Nat>, Nat, values => Shape.offset(left, Shape.project_indices(left, values)), Shape.indices(output, coordinate), ds, Coordinates.digits_unflatten(output, index)) : {_ == Shape.row_major(left, projected) : Nat}
147 %Equal.cong(List<&2, Nat>, Nat, values => Shape.offset(left, values), Shape.indices(left, projected), Shape.project_indices(left, Shape.indices(output, coordinate)), Broadcast.project_left(left, right, output, alignment, coordinate)) : {_ == Shape.row_major(left, projected) : Nat}
148 Coordinates.offset_indices(left, projected)
149
150def walk_right_value(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +alignment: Shape.Broadcast<left, right, output>, +index: Shape.Fin(Shape.numel(output))) -> {Shape.walk_right(Shape.index_table(left, right, output), Shape.Fin.value(Shape.numel(output), index)) == Shape.row_major(right, Shape.project_right(left, right, output, alignment, Shape.unflatten(output, index))) : Nat}:
151 +value = Shape.Fin.value(Shape.numel(output), index)
152 +ds = Shape.digits(output, value)
153 +nl = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, left))
154 +nr = Nat.sub(List.length(&2, Nat, output), List.length(&2, Nat, right))
155 +coordinate = Shape.unflatten(output, index)
156 +projected = Shape.project_right(left, right, output, alignment, coordinate)
157 +strides = strides_of(right)
158 +dropped = List.drop(&2, Nat, ds, nr)
159 %Equal.sym(Nat, Shape.walk_right(Shape.index_rows(Shape.pad(nl, left), Shape.pad(nr, right), output), value), dot(strides_of(Shape.pad(nr, right)), ds), walk_right_rows(Shape.pad(nl, left), Shape.pad(nr, right), output, value, left_padded(left, right, output, alignment))) : {_ == Shape.row_major(right, projected) : Nat}
160 %Equal.cong(List<&2, Nat>, Nat, steps => dot(steps, ds), zeros(nr, strides), strides_of(Shape.pad(nr, right)), Equal.sym(List<&2, Nat>, strides_of(Shape.pad(nr, right)), zeros(nr, strides), pad_strides(nr, right))) : {_ == Shape.row_major(right, projected) : Nat}
161 %Equal.sym(Nat, dot(zeros(nr, strides), ds), dot(strides, dropped), dot_zeros(nr, ds, strides)) : {_ == Shape.row_major(right, projected) : Nat}
162 %Equal.sym(Nat, dot(strides, dropped), Shape.offset(right, Shape.broadcast_indices(right, dropped)), dot_broadcast(right, dropped)) : {_ == Shape.row_major(right, projected) : Nat}
163 %Equal.cong(Nat, Nat, width => Shape.offset(right, Shape.broadcast_indices(right, List.drop(&2, Nat, ds, Nat.sub(width, List.length(&2, Nat, right))))), List.length(&2, Nat, ds), List.length(&2, Nat, output), digits_length(output, value)) : {_ == Shape.row_major(right, projected) : Nat}
164 %Equal.cong(List<&2, Nat>, Nat, values => Shape.offset(right, Shape.project_indices(right, values)), Shape.indices(output, coordinate), ds, Coordinates.digits_unflatten(output, index)) : {_ == Shape.row_major(right, projected) : Nat}
165 %Equal.cong(List<&2, Nat>, Nat, values => Shape.offset(right, values), Shape.indices(right, projected), Shape.project_indices(right, Shape.indices(output, coordinate)), Broadcast.project_right(left, right, output, alignment, coordinate)) : {_ == Shape.row_major(right, projected) : Nat}
166 Coordinates.offset_indices(right, projected)