def stretch

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul.bend:237

import Base

237# Index `index` of a batch axis of `extent`: the index itself, or 0 where the
238# extent is 1 and stretches.
239def stretch(+index: Nat, one: Bool) -> Nat:
240  match one:
241    case True{}:
242      0n
243    case False{}:
244      index