def operand_ids

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward.bend:1072

import Base

1072# The operand ids of a node, in operand order.
1073def operand_ids(node: Node) -> List<&2, Nat>:
1074  match node:
1075    case Leaf{}:
1076      []
1077    case Add{a, b, left, right}:
1078      [a, b]
1079    case Sub{a, b, left, right}:
1080      [a, b]
1081    case Mul{a, b, left, right}:
1082      [a, b]
1083    case Sum{a, input}:
1084      [a]
1085    case SumDims{a, input, kept}:
1086      [a]
1087    case Matmul{a, b, left, right}:
1088      [a, b]
1089    case MatmulNT{a, b, left, right}:
1090      [a, b]
1091    case Relu{a, input}:
1092      [a]
1093    case MseLoss{a, b, left, right}:
1094      [a, b]
1095    case CrossEntropy{a, input}:
1096      [a]
1097    case BatchedMatmul{a, b, left, right}:
1098      [a, b]
1099    case Slice{a, input, kept, plan}:
1100      [a]
1101    case Gathered{a, input, plan, entries}:
1102      [a]
1103    case Constant{}:
1104      []
1105    case Frozen{id, inner}:
1106      operand_ids(inner)