def targets

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward_LAWS.bend:37

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

37# Operand bounds, stated without the engine's validator: a node may name only
38# nodes older than itself, so each operand id must be below the number of
39# older nodes.
40def targets(node: Backward.Node) -> List<&2, Nat>:
41  match node:
42    case Backward.Leaf{}:
43      []
44    case Backward.Add{a, b, left, right}:
45      [a, b]
46    case Backward.Sub{a, b, left, right}:
47      [a, b]
48    case Backward.Mul{a, b, left, right}:
49      [a, b]
50    case Backward.Sum{a, input}:
51      [a]
52    case Backward.SumDims{a, input, kept}:
53      [a]
54    case Backward.Matmul{a, b, left, right}:
55      [a, b]
56    case Backward.MatmulNT{a, b, left, right}:
57      [a, b]
58    case Backward.Relu{a, input}:
59      [a]
60    case Backward.MseLoss{a, b, left, right}:
61      [a, b]
62    case Backward.CrossEntropy{a, input}:
63      [a]
64    case Backward.BatchedMatmul{a, b, left, right}:
65      [a, b]
66    case Backward.Slice{a, input, kept, plan}:
67      [a]
68    case Backward.Gathered{a, input, plan, entries}:
69      [a]
70    case Backward.Constant{}:
71      []
72    case Backward.Frozen{id, inner}:
73      targets(inner)