def take

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

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

62# Takes out of a newest-first log every contribution to id and accumulates
63# them into slot, oldest first; the others stay in the log, in order.
64def take(log: Log(), +id: Nat, slot: Gradient.Slot) -> Taken():
65  match log:
66    case Nil{}:
67      (Done{slot}, Nil{})
68    case Con{Backward.Contribution{+target, +shape, gradient}, rest}:
69      sorted(target, shape, gradient, Nat.is_eq(id, target), take(rest, id, slot))