def splits

bendtorch-core@0.1.0.0 · src/autograd/reduce.bend:241

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix

241# How many top levels of the target tree fork. Every fork clones the whole
242# contribution. A target whose fibers are contiguous splits by work, as
243# `Matrix.splits` does with the terms of a cell; a target with strided fibers
244# never forks, since there the forked build ran slower at 8 and 24 threads
245# than one thread. The proofs hold for any split.
246def splits(+plan: Shape.BroadcastPlan, +depth: Nat, +length: Nat) -> Nat:
247  +extents = Shape.extents(Shape.output_shape(plan))
248  +rows = Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), extents)
249  Bool.pick(Nat, contiguous(rows, extents), Matrix.splits(depth, terms(plan, length)), 0n)