def contiguous

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

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

229# Whether every fiber is one contiguous run of the contribution: no axis the
230# target keeps comes after an axis it reduces, counting only axes of extent
231# above 1. Row sums such as [n, m] to [n, 1] are contiguous; column and batch
232# sums such as [n, m] to [1, m] or [b, n, m] to [n, m] are strided.
233def contiguous(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>) -> Bool:
234  match rows extents:
235    case Shape.Row{stride, +left, right} <> +rest +extent <> +others:
236      Bool.pick(Bool, Bool.and(Nat.is_lt(1n, extent), Nat.is_eq(left, 0n)), Bool.not(kept_after(rest, others)),
237        contiguous(rest, others))
238    case _ _:
239      True{}