src/autograd/sgd_PROOF.bend
bendtorch-core@0.1.0.0 · 172 lines · 1 definition
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
5import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_proofs.bend as BinaryProof
6
7import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_PROOF.bend as BinaryLaw
8import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
9import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
10import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as TensorBackward
11import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_PROOF.bend as TensorBackwardProof
12import bendtorch-core@0.1.0.0/src/autograd/sgd.bend as Sgd
13import bendtorch-core@0.1.0.0/src/autograd/sgd_LAWS.bend as Laws
14
15def Laws.sgd_descend_scalar(lr, p, g):
16 {==}
17
18def Laws.sgd_empty_slot_kept(lr, node, shape, value, saved):
19 {==}
20
21def Laws.sgd_step_values(lr, node, shape, saved, g_shape, plan, planned, fill, ops, fits, g_fill, g_ops, g_fits, P, h):
22 +planned = planned
23 BinaryProof.binary_replayed(Elementwise.Descend{lr}, plan, shape, fill, ops, g_shape, g_fill, g_ops, planned,
24 pair => P(Sgd.descended(node, shape, saved, pair)), h)(fits, g_fits)
25
26def Laws.sgd_update_reaches_target(lr, newer, entry, rest, older, fits):
27 match newer:
28 case Nil{}:
29 %TensorBackwardProof.add_zero(older) :
30 {Sgd.descend_at(lr, entry <> rest, 1n+_, older, Backward.at_head(1n+_, older)) ==
31 Sgd.placed(rest, Sgd.descend(lr, entry)) : Sgd.Walked()}
32 %TensorBackwardProof.eq_refl(older) :
33 {Sgd.descend_at(lr, entry <> rest, 1n+older, older, _) == Sgd.placed(rest, Sgd.descend(lr, entry)) : Sgd.Walked()}
34 {==}
35 case Con{head, tail}:
36 %TensorBackwardProof.older_ne(older, List.length(&1, Backward.Entry, tail)) :
37 {Sgd.descend_at(lr, head <> List.append(&1, Backward.Entry, tail, entry <> rest),
38 1n+Nat.add(older, 1n+List.length(&1, Backward.Entry, tail)), older, _) ==
39 Sgd.kept(head, Laws.newer_kept(tail, Sgd.placed(rest, Sgd.descend(lr, entry)))) : Sgd.Walked()}
40 %TensorBackwardProof.add_succ(older, List.length(&1, Backward.Entry, tail)) :
41 {Sgd.descend_at(lr, head <> List.append(&1, Backward.Entry, tail, entry <> rest), 1n+_, older, False{}) ==
42 Sgd.kept(head, Laws.newer_kept(tail, Sgd.placed(rest, Sgd.descend(lr, entry)))) : Sgd.Walked()}
43 Equal.cong(Sgd.Walked(), Sgd.Walked(), walked => Sgd.kept(head, walked),
44 Sgd.descend_at(lr, List.append(&1, Backward.Entry, tail, entry <> rest), 1n+Nat.add(older, List.length(&1, Backward.Entry, tail)),
45 older, Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, tail)), older)),
46 Laws.newer_kept(tail, Sgd.placed(rest, Sgd.descend(lr, entry))),
47 Laws.sgd_update_reaches_target(lr, tail, entry, rest, older, fits))
48
49def Laws.sgd_only_leaves_accepted(node):
50 match node:
51 case Backward.Leaf{}:
52 {==}
53 case Backward.Add{a, b, left, right}:
54 {==}
55 case Backward.Sub{a, b, left, right}:
56 {==}
57 case Backward.Mul{a, b, left, right}:
58 {==}
59 case Backward.Sum{a, input}:
60 {==}
61 case Backward.SumDims{a, input, kept}:
62 {==}
63 case Backward.Matmul{a, b, left, right}:
64 {==}
65 case Backward.MatmulNT{a, b, left, right}:
66 {==}
67 case Backward.Relu{a, input}:
68 {==}
69 case Backward.MseLoss{a, b, left, right}:
70 {==}
71 case Backward.CrossEntropy{a, input}:
72 {==}
73 case Backward.BatchedMatmul{a, b, left, right}:
74 {==}
75 case Backward.Slice{a, input, kept, plan}:
76 {==}
77 case Backward.Gathered{a, input, plan, entries}:
78 {==}
79 case Backward.Constant{}:
80 {==}
81 case Backward.Frozen{id, inner}:
82 Laws.sgd_only_leaves_accepted(inner)
83
84def Laws.sgd_no_params(lr, count, entries):
85 {==}
86
87def Laws.sgd_params_in_order(lr, count, id, rest, entries):
88 {==}
89
90def Laws.sgd_refusal_reads_target(newer, node, rest, older, fits):
91 match newer:
92 case Nil{}:
93 %TensorBackwardProof.add_zero(older) :
94 {Sgd.refusal_at(node <> rest, 1n+_, older, Backward.at_head(1n+_, older)) == Sgd.leaf_refusal(node) : Maybe<&2, Sgd.Error>}
95 %TensorBackwardProof.eq_refl(older) :
96 {Sgd.refusal_at(node <> rest, 1n+older, older, _) == Sgd.leaf_refusal(node) : Maybe<&2, Sgd.Error>}
97 {==}
98 case Con{head, tail}:
99 %TensorBackwardProof.older_ne(older, List.length(&2, Backward.Node, tail)) :
100 {Sgd.refusal_at(head <> List.append(&2, Backward.Node, tail, node <> rest),
101 1n+Nat.add(older, 1n+List.length(&2, Backward.Node, tail)), older, _) == Sgd.leaf_refusal(node) : Maybe<&2, Sgd.Error>}
102 %TensorBackwardProof.add_succ(older, List.length(&2, Backward.Node, tail)) :
103 {Sgd.refusal_at(head <> List.append(&2, Backward.Node, tail, node <> rest), 1n+_, older, False{}) ==
104 Sgd.leaf_refusal(node) : Maybe<&2, Sgd.Error>}
105 Laws.sgd_refusal_reads_target(tail, node, rest, older, fits)
106
107
108def ne_succ(+number: Nat, -extra: Nat) -> {False{} == Nat.is_eq(number, 1n+Nat.add(number, extra)) : Bool}:
109 match number:
110 case 0n:
111 {==}
112 case 1n+rest:
113 ne_succ(rest, extra)
114
115def Laws.sgd_refusal_unknown(nodes, extra):
116 match nodes:
117 case Nil{}:
118 {==}
119 case Con{node, +rest}:
120 +below = List.length(&2, Backward.Node, rest)
121 %ne_succ(below, extra) :
122 {Sgd.refusal_at(node <> rest, 1n+below, 1n+Nat.add(below, extra), _) == Some{Sgd.UnknownNode{}} : Maybe<&2, Sgd.Error>}
123 %Equal.sym(Nat, 1n+Nat.add(below, extra), Nat.add(below, 1n+extra), TensorBackwardProof.add_succ(below, extra)) :
124 {Sgd.refusal_at(rest, below, _, Backward.at_head(below, _)) == Some{Sgd.UnknownNode{}} : Maybe<&2, Sgd.Error>}
125 Laws.sgd_refusal_unknown(rest, 1n+extra)
126
127def Laws.sgd_step_refused(entries, lr, params, error, refused):
128 %TensorBackwardProof.listed_back(entries) :
129 {Sgd.planned(lr, params, _) == (Backward.Tape{entries}, Fail{error}) : Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
130 %Equal.sym(Maybe<&2, Sgd.Error>, Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params), Some{error}, refused) :
131 {Sgd.checked(lr, entries, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, _) ==
132 (Backward.Tape{entries}, Fail{error}) : Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
133 {==}
134
135def Laws.sgd_negative_rate_refused(lr, below):
136 %Equal.sym(Bool, Elementwise.below_zero(lr), True{}, below) :
137 {Sgd.refused_if(_) == Some{Sgd.InvalidLearningRate{}} : Maybe<&2, Sgd.Error>}
138 {==}
139
140def Laws.sgd_rate_accepted(lr, above):
141 %Equal.sym(Bool, Elementwise.below_zero(lr), False{}, above) :
142 {Sgd.refused_if(_) == None{} : Maybe<&2, Sgd.Error>}
143 {==}
144
145def Laws.sgd_rate_edges():
146 {==}
147
148def Laws.sgd_step_rate_refused(entries, lr, params, error, accepted, refused):
149 %TensorBackwardProof.listed_back(entries) :
150 {Sgd.planned(lr, params, _) == (Backward.Tape{entries}, Fail{error}) : Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
151 %Equal.sym(Maybe<&2, Sgd.Error>, Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params), None{}, accepted) :
152 {Sgd.checked(lr, entries, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, _) ==
153 (Backward.Tape{entries}, Fail{error}) : Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
154 %Equal.sym(Maybe<&2, Sgd.Error>, Sgd.rate_refusal(lr), Some{error}, refused) :
155 {Sgd.rated(lr, entries, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, _) ==
156 (Backward.Tape{entries}, Fail{error}) : Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
157 {==}
158
159def Laws.sgd_step_accepted(entries, lr, params, accepted, rated):
160 %TensorBackwardProof.listed_back(entries) :
161 {Sgd.planned(lr, params, _) ==
162 Sgd.walked_tape(Sgd.descend_all(lr, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, entries)) :
163 Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
164 %Equal.sym(Maybe<&2, Sgd.Error>, Sgd.refusal(TensorBackward.nodes_of(entries), List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params), None{}, accepted) :
165 {Sgd.checked(lr, entries, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, _) ==
166 Sgd.walked_tape(Sgd.descend_all(lr, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, entries)) :
167 Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
168 %Equal.sym(Maybe<&2, Sgd.Error>, Sgd.rate_refusal(lr), None{}, rated) :
169 {Sgd.rated(lr, entries, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, _) ==
170 Sgd.walked_tape(Sgd.descend_all(lr, List.length(&2, Backward.Node, TensorBackward.nodes_of(entries)), params, entries)) :
171 Backward.Tape & Result<&2, &2, Sgd.Error, Unit>}
172 {==}