def chained

bendtorch-core@0.1.0.0 · src/autograd/tensor_chain_LAWS.bend:114

import Base
import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward

114# The working gradients by the chain rule, pulled, for entries newest first
115# that start with slots and receive the contributions in log: the entry with
116# id k gathers every contribution sent to k into its starting slot; if it then
117# holds g, the contributions of vjp(node, saved, g) join the log.
118def chained(entries: List<&1, Backward.Entry>, slots: List<&1, Gradient.Slot>, +count: Nat, log: Log()) -> Slots():
119  match entries slots count:
120    case Nil{} slots count:
121      gather(slots, count, log)
122    case Con{entry, rest} Nil{} count:
123      Done{Nil{}}
124    case Con{entry, rest} Con{slot, more} 0n:
125      Done{Nil{}}
126    case Con{Backward.Entry{node, shape, value, saved, grad}, rest} Con{slot, more} 1n+below:
127      reached(node, saved, take(log, below, slot), next => chained(rest, more, below, next))