1	# bend-ml-autograd: diferenciação automática com uma lei provada.
2	#
3	#   import bend-ml-autograd@0.1.0.0/main.bend as AG
4	#
5	# 1) MODELO PROVADO (sobre Nat): o modo reverso do autodiff dá o mesmo resultado
6	#    que o modo direto (números duais). Prova pelo kernel, usando os lemas de
7	#    soma e produto do bend-ml-nat-lemmas.
8	# 2) AUTOGRAD DE ESCALARES em F32 (estilo micrograd), com a mesma estrutura.
9	# 3) CAMADAS COM TENSORES (forward e backward) cujo formato é checado pelo tipo.
10	#
11	# F32 não é um número real, então a correção NUMÉRICA do F32 é validada por
12	# testes contra o PyTorch (reference/test_autograd.py), não por prova.
13	
14	import Base
15	import bend-ml-nat-lemmas@0.1.0.0/main.bend as NL
16	import bend-ml-tensor@0.1.0.0/main.bend as T
17	
18	# =====================================================================
19	# 1) Modelo provado
20	# =====================================================================
21	
22	# Expressões com uma variável NX, sobre Nat (um semianel comutativo).
23	type NE is Data:
24	  NCst{n: Nat}
25	  NX{}
26	  NAdd{a: NE, b: NE}
27	  NMul{a: NE, b: NE}
28	
29	# valor da expressão em x
30	def nval(e: NE, +x: Nat) -> Nat:
31	  match e:
32	    case NCst{n}:
33	      n
34	    case NX{}:
35	      x
36	    case NAdd{a, b}:
37	      Nat.add(nval(a, x), nval(b, x))
38	    case NMul{+a, +b}:
39	      Nat.mul(nval(a, x), nval(b, x))
40	
41	# derivada em modo DIRETO (números duais): derivada da soma e regra do produto
42	def nfwd(e: NE, +x: Nat) -> Nat:
43	  match e:
44	    case NCst{n}:
45	      0n
46	    case NX{}:
47	      1n
48	    case NAdd{a, b}:
49	      Nat.add(nfwd(a, x), nfwd(b, x))
50	    case NMul{+a, +b}:
51	      Nat.add(Nat.mul(nfwd(a, x), nval(b, x)), Nat.mul(nval(a, x), nfwd(b, x)))
52	
53	# derivada em modo REVERSO: g é o gradiente que chega da saída; cada nó repassa
54	# g (soma) ou g * valor-do-outro-fator (produto) para os filhos e soma o que
55	# chega em NX.
56	def nbwd(e: NE, +x: Nat, +g: Nat) -> Nat:
57	  match e:
58	    case NCst{n}:
59	      0n
60	    case NX{}:
61	      g
62	    case NAdd{a, b}:
63	      Nat.add(nbwd(a, x, g), nbwd(b, x, g))
64	    case NMul{+a, +b}:
65	      Nat.add(nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))))
66	
67	# a * (b + c) == a*b + a*c   (distributividade à esquerda; o nat-lemmas só tem à direita)
68	def ndist_l(+a: Nat, +b: Nat, +c: Nat) -> {Nat.mul(a, Nat.add(b, c)) == Nat.add(Nat.mul(a, b), Nat.mul(a, c)) : Nat}:
69	  Equal.trans(Nat, Nat.mul(a, Nat.add(b, c)), Nat.mul(Nat.add(b, c), a), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
70	    NL.mul_comm(a, Nat.add(b, c)),
71	    Equal.trans(Nat, Nat.mul(Nat.add(b, c), a), Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
72	      Equal.sym(Nat, Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.mul(Nat.add(b, c), a), NL.mul_dist(b, c, a)),
73	      Equal.trans(Nat, Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
74	        Equal.cong(Nat, Nat, k => Nat.add(k, Nat.mul(c, a)), Nat.mul(b, a), Nat.mul(a, b), NL.mul_comm(b, a)),
75	        Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(a, b), k), Nat.mul(c, a), Nat.mul(a, c), NL.mul_comm(c, a)))))
76	
77	# Caso do produto, com tudo abstraído em números:
78	# (g*vb)*fa + (g*va)*fb == g * (fa*vb + va*fb)
79	def nmul_case(+g: Nat, +fa: Nat, +fb: Nat, +va: Nat, +vb: Nat, +ba: Nat, +bb: Nat, iha: {ba == Nat.mul(Nat.mul(g, vb), fa) : Nat}, ihb: {bb == Nat.mul(Nat.mul(g, va), fb) : Nat}) -> {Nat.add(ba, bb) == Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))) : Nat}:
80	  Equal.trans(Nat, Nat.add(ba, bb), Nat.add(Nat.mul(Nat.mul(g, vb), fa), bb), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
81	    Equal.cong(Nat, Nat, k => Nat.add(k, bb), ba, Nat.mul(Nat.mul(g, vb), fa), iha),
82	    Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(g, vb), fa), bb), Nat.add(Nat.mul(Nat.mul(g, vb), fa), Nat.mul(Nat.mul(g, va), fb)), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
83	      Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(Nat.mul(g, vb), fa), k), bb, Nat.mul(Nat.mul(g, va), fb), ihb),
84	      Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(g, vb), fa), Nat.mul(Nat.mul(g, va), fb)), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(Nat.mul(g, va), fb)), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
85	        Equal.cong(Nat, Nat, k => Nat.add(k, Nat.mul(Nat.mul(g, va), fb)), Nat.mul(Nat.mul(g, vb), fa), Nat.mul(g, Nat.mul(fa, vb)),
86	          Equal.trans(Nat, Nat.mul(Nat.mul(g, vb), fa), Nat.mul(g, Nat.mul(vb, fa)), Nat.mul(g, Nat.mul(fa, vb)),
87	            Equal.sym(Nat, Nat.mul(g, Nat.mul(vb, fa)), Nat.mul(Nat.mul(g, vb), fa), NL.mul_assoc(g, vb, fa)),
88	            Equal.cong(Nat, Nat, k => Nat.mul(g, k), Nat.mul(vb, fa), Nat.mul(fa, vb), NL.mul_comm(vb, fa)))),
89	        Equal.trans(Nat, Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(Nat.mul(g, va), fb)), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(g, Nat.mul(va, fb))), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
90	          Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(g, Nat.mul(fa, vb)), k), Nat.mul(Nat.mul(g, va), fb), Nat.mul(g, Nat.mul(va, fb)),
91	            Equal.sym(Nat, Nat.mul(g, Nat.mul(va, fb)), Nat.mul(Nat.mul(g, va), fb), NL.mul_assoc(g, va, fb))),
92	          Equal.sym(Nat, Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(g, Nat.mul(va, fb))), ndist_l(g, Nat.mul(fa, vb), Nat.mul(va, fb)))))))
93	
94	# O modo reverso devolve g vezes a derivada do modo direto.
95	def nbwd_ok(e: NE, +x: Nat, +g: Nat) -> {nbwd(e, x, g) == Nat.mul(g, nfwd(e, x)) : Nat}:
96	  match e:
97	    case NCst{n}:
98	      Equal.sym(Nat, Nat.mul(g, 0n), 0n, NL.mul_zero(g))
99	    case NX{}:
100	      Equal.sym(Nat, Nat.mul(g, 1n), g, NL.mul_one_r(g))
101	    case NAdd{+a, +b}:
102	      Equal.trans(Nat, Nat.add(nbwd(a, x, g), nbwd(b, x, g)), Nat.add(Nat.mul(g, nfwd(a, x)), nbwd(b, x, g)), Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))),
103	        Equal.cong(Nat, Nat, k => Nat.add(k, nbwd(b, x, g)), nbwd(a, x, g), Nat.mul(g, nfwd(a, x)), nbwd_ok(a, x, g)),
104	        Equal.trans(Nat, Nat.add(Nat.mul(g, nfwd(a, x)), nbwd(b, x, g)), Nat.add(Nat.mul(g, nfwd(a, x)), Nat.mul(g, nfwd(b, x))), Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))),
105	          Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(g, nfwd(a, x)), k), nbwd(b, x, g), Nat.mul(g, nfwd(b, x)), nbwd_ok(b, x, g)),
106	          Equal.sym(Nat, Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))), Nat.add(Nat.mul(g, nfwd(a, x)), Nat.mul(g, nfwd(b, x))), ndist_l(g, nfwd(a, x), nfwd(b, x)))))
107	    case NMul{+a, +b}:
108	      nmul_case(g, nfwd(a, x), nfwd(b, x), nval(a, x), nval(b, x), nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))), nbwd_ok(a, x, Nat.mul(g, nval(b, x))), nbwd_ok(b, x, Nat.mul(g, nval(a, x))))
109	
110	# LEI (reverse_eq_forward): para qualquer expressão feita de constantes, X, somas
111	# e produtos, o gradiente em modo reverso (começando com gradiente 1 na saída)
112	# é igual à derivada em modo direto. Vale para todo x.
113	law reverse_eq_forward:
114	  for +e: NE
115	  for +x: Nat
116	  {nbwd(e, x, 1n) == nfwd(e, x) : Nat}
117	
118	def reverse_eq_forward(e, x):
119	  Equal.trans(Nat, nbwd(e, x, 1n), Nat.mul(1n, nfwd(e, x)), nfwd(e, x), nbwd_ok(e, x, 1n), NL.mul_one_l(nfwd(e, x)))
120	
121	
122	# =====================================================================
123	# 2) Autograd de escalares em F32 (reverse-mode, estilo micrograd)
124	# =====================================================================
125	
126	type G is Data:
127	  GCst{c: F32}
128	  GVar{i: Nat}
129	  GAdd{a: G, b: G}
130	  GMul{a: G, b: G}
131	  GRelu{a: G}
132	  GTanh{a: G}
133	  GExp{a: G}
134	
135	def lookup(env: List<&2, F32>, +i: Nat) -> F32:
136	  match env i:
137	    case Nil{} _:
138	      0.0
139	    case Con{h, t} 0n:
140	      h
141	    case Con{h, t} 1n+p:
142	      lookup(t, p)
143	
144	def gval(e: G, +env: List<&2, F32>) -> F32:
145	  match e:
146	    case GCst{c}:
147	      c
148	    case GVar{i}:
149	      lookup(env, i)
150	    case GAdd{a, b}:
151	      (gval(a, env) + gval(b, env) : F32)
152	    case GMul{a, b}:
153	      (gval(a, env) * gval(b, env) : F32)
154	    case GRelu{a}:
155	      F32.max(gval(a, env), 0.0)
156	    case GTanh{a}:
157	      F32.tanh(gval(a, env))
158	    case GExp{a}:
159	      F32.exp(gval(a, env))
160	
161	# soma v ao gradiente da variável i
162	def acc_add(acc: List<&2, F32>, +i: Nat, +v: F32) -> List<&2, F32>:
163	  match acc i:
164	    case Nil{} _:
165	      Nil{}
166	    case Con{+h, t} 0n:
167	      (h + v : F32) <> t
168	    case Con{h, t} 1n+p:
169	      h <> acc_add(t, p, v)
170	
171	def mask2(pos: Bool) -> F32:
172	  match pos:
173	    case True{}:
174	      1.0
175	    case False{}:
176	      0.0
177	
178	# derivada local da relu: 1 se a entrada > 0, senão 0
179	def relu_d(+x: F32) -> F32:
180	  mask2(F32.is_gt(x, 0.0))
181	
182	def tanh_d(+x: F32) -> F32:
183	  (1.0 - (F32.tanh(x) * F32.tanh(x) : F32) : F32)
184	
185	# g é o gradiente que chega ao nó; acc acumula o gradiente de cada variável
186	def gbwd(e: G, +env: List<&2, F32>, +g: F32, acc: List<&2, F32>) -> List<&2, F32>:
187	  match e:
188	    case GCst{c}:
189	      acc
190	    case GVar{+i}:
191	      acc_add(acc, i, g)
192	    case GAdd{a, b}:
193	      gbwd(b, env, g, gbwd(a, env, g, acc))
194	    case GMul{+a, +b}:
195	      gbwd(b, env, (g * gval(a, env) : F32), gbwd(a, env, (g * gval(b, env) : F32), acc))
196	    case GRelu{+a}:
197	      gbwd(a, env, (g * relu_d(gval(a, env)) : F32), acc)
198	    case GTanh{+a}:
199	      gbwd(a, env, (g * tanh_d(gval(a, env)) : F32), acc)
200	    case GExp{+a}:
201	      gbwd(a, env, (g * F32.exp(gval(a, env)) : F32), acc)
202	
203	# gradiente da expressão em relação a cada variável de env
204	def ggrad(e: G, +env: List<&2, F32>) -> List<&2, F32>:
205	  gbwd(e, env, 1.0, T.fill.l(List.length(&2, F32, env), 0.0))
206	
207	# =====================================================================
208	# 3) Camadas com tensores: o formato de cada gradiente é imposto pelo TIPO
209	# =====================================================================
210	
211	# y = x·W + b           x: n x i,  W: i x o,  b: o   ->   y: n x o
212	def linear(-n: Nat, +i: Nat, +o: Nat, x: T.Mat<n, i>, w: T.Mat<i, o>, b: T.Vec<o>) -> T.Mat<n, o>:
213	  T.Mat.add_row(n, o, T.Mat.matmul(n, i, o, x, w), b)
214	
215	# Dado o gradiente dy que chega em y, devolve (dx, (dW, db)):
216	#   dx = dy · Wᵀ   (n x o)·(o x i) = n x i
217	#   dW = xᵀ · dy   (i x n)·(n x o) = i x o
218	#   db = soma das linhas de dy
219	# Se alguma conta tivesse as dimensões erradas, o programa não compilaria.
220	def linear_bwd(+n: Nat, +i: Nat, +o: Nat, x: T.Mat<n, i>, w: T.Mat<i, o>, +dy: T.Mat<n, o>) -> T.Mat<n, i> & (T.Mat<i, o> & T.Vec<o>):
221	  (T.Mat.matmul(n, o, i, dy, T.Mat.transpose(i, o, w)), (T.Mat.matmul(i, n, o, T.Mat.transpose(n, i, x), dy), T.Mat.col_sums(n, o, dy)))
222	
223	def mask.l(xs: List<&2, F32>) -> List<&2, F32>:
224	  match xs:
225	    case Nil{}:
226	      Nil{}
227	    case Con{h, t}:
228	      relu_d(h) <> mask.l(t)
229	
230	def mask.rows(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
231	  match xs:
232	    case Nil{}:
233	      Nil{}
234	    case Con{r, t}:
235	      mask.l(r) <> mask.rows(t)
236	
237	# y = relu(x); o gradiente passa onde x > 0
238	def relu_bwd(-r: Nat, -c: Nat, x: T.Mat<r, c>, dy: T.Mat<r, c>) -> T.Mat<r, c>:
239	  match x:
240	    case T.Mat{rows}:
241	      T.Mat.mul(r, c, dy, T.Mat{mask.rows(rows)})
242	
243	# one-hot: n linhas de c números, com 1 na posição do rótulo
244	def onehot.cell(hit: Bool) -> F32:
245	  mask2(hit)
246	
247	def onehot.row(c: Nat, +k: Nat, +j: Nat) -> List<&2, F32>:
248	  match c:
249	    case 0n:
250	      Nil{}
251	    case 1n+p:
252	      onehot.cell(Nat.is_eq(k, j)) <> onehot.row(p, k, 1n+j)
253	
254	def onehot.rows(labels: List<&2, Nat>, +c: Nat) -> List<&2, List<&2, F32>>:
255	  match labels:
256	    case Nil{}:
257	      Nil{}
258	    case Con{k, t}:
259	      onehot.row(c, k, 0n) <> onehot.rows(t, c)
260	
261	def log.l(xs: List<&2, F32>) -> List<&2, F32>:
262	  match xs:
263	    case Nil{}:
264	      Nil{}
265	    case Con{+h, t}:
266	      F32.log(F32.max(h, 0.000000000001)) <> log.l(t)
267	
268	def log.rows(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
269	  match xs:
270	    case Nil{}:
271	      Nil{}
272	    case Con{r, t}:
273	      log.l(r) <> log.rows(t)
274	
275	def rows.total(xs: List<&2, List<&2, F32>>, acc: F32) -> F32:
276	  match xs:
277	    case Nil{}:
278	      acc
279	    case Con{r, t}:
280	      rows.total(t, (acc + T.sum.go(r, 0.0) : F32))
281	
282	def ce_loss2(-n: Nat, -c: Nat, +cc: Nat, +nf: F32, p: T.Mat<n, c>, labels: List<&2, Nat>) -> F32:
283	  match p:
284	    case T.Mat{rows}:
285	      (0.0 - (rows.total(T.rows.zip_mul(log.rows(rows), onehot.rows(labels, cc)), 0.0) / nf : F32) : F32)
286	
287	# perda de entropia cruzada média: -(1/n) Σ log p[i][rótulo_i]
288	def ce_loss(+n: Nat, +c: Nat, logits: T.Mat<n, c>, labels: List<&2, Nat>) -> F32:
289	  ce_loss2(n, c, c, F32.from_nat(n), T.Mat.softmax(n, c, logits), labels)
290	
291	# gradiente da perda em relação aos logits: (softmax - one-hot) / n
292	def ce_grad(+n: Nat, +c: Nat, logits: T.Mat<n, c>, labels: List<&2, Nat>) -> T.Mat<n, c>:
293	  T.Mat.scale(n, c, (1.0 / F32.from_nat(n) : F32), T.Mat.sub(n, c, T.Mat.softmax(n, c, logits), T.Mat{onehot.rows(labels, c)}))
294	
295	def sgd(-r: Nat, -c: Nat, +lr: F32, w: T.Mat<r, c>, dw: T.Mat<r, c>) -> T.Mat<r, c>:
296	  T.Mat.sub(r, c, w, T.Mat.scale(r, c, lr, dw))
297	
298	def sgd_vec(-n: Nat, +lr: F32, b: T.Vec<n>, db: T.Vec<n>) -> T.Vec<n>:
299	  T.Vec.sub(n, b, T.Vec.scale(n, lr, db))