def gbwd

bend-ml-autograd@0.1.1.0 · main.bend:185

import Base

185# g is the gradient arriving at the node; acc accumulates the gradient of each variable
186def 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)