1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise.bend as Elementwise
5	import ../tensor/contracts/binary_proofs.bend as BinaryProof
6	# Proves binary_replayed, which importing the proof above would otherwise leave open.
7	import ../tensor/contracts/binary_PROOF.bend as BinaryLaw
8	import ./gradient.bend as Gradient
9	import ./tensor_backward.bend as Backward
10	import ./tensor_backward_LAWS.bend as TensorBackward
11	import ./tensor_backward_PROOF.bend as TensorBackwardProof
12	import ./sgd.bend as Sgd
13	import ./sgd_LAWS.bend as Laws
14	
15	def Laws.sgd_descend_scalar(lr, p, g):
16	  {==}
17	
18	def Laws.sgd_empty_slot_kept(lr, node, shape, value, saved):
19	  {==}
20	
21	def 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	
26	def 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	
49	def 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	
84	def Laws.sgd_no_params(lr, count, entries):
85	  {==}
86	
87	def Laws.sgd_params_in_order(lr, count, id, rest, entries):
88	  {==}
89	
90	def 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	# A list of n nodes never has id n + 1 + extra at its head.
108	def 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	
115	def 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	
127	def 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	
135	def 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	
140	def 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	
145	def Laws.sgd_rate_edges():
146	  {==}
147	
148	def 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	
159	def 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	  {==}