def operands_below

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

import Base

823def operands_below(node: Node, +bound: Nat) -> Bool:
824  match node:
825    case Leaf{}:
826      True{}
827    case Add{a, b, left, right}:
828      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
829    case Sub{a, b, left, right}:
830      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
831    case Mul{a, b, left, right}:
832      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
833    case Sum{a, input}:
834      Nat.is_lt(a, bound)
835    case SumDims{a, input, kept}:
836      Nat.is_lt(a, bound)
837    case Matmul{a, b, left, right}:
838      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
839    case MatmulNT{a, b, left, right}:
840      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
841    case Relu{a, input}:
842      Nat.is_lt(a, bound)
843    case MseLoss{a, b, left, right}:
844      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
845    case CrossEntropy{a, input}:
846      Nat.is_lt(a, bound)
847    case BatchedMatmul{a, b, left, right}:
848      Bool.and(Nat.is_lt(a, bound), Nat.is_lt(b, bound))
849    case Slice{a, input, kept, plan}:
850      Nat.is_lt(a, bound)
851    case Gathered{a, input, plan, entries}:
852      Nat.is_lt(a, bound)
853    case Constant{}:
854      True{}
855    case Frozen{id, inner}:
856      operands_below(inner, bound)