def stretch_spec

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_PROOF.bend:343

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

343# The batch number of the build reads each operand's batch as the shape
344# module broadcasts it.
345# An index of a batch axis: the stretch of the build, the shape module's
346# broadcast index.
347def stretch_spec(+extent: Nat, +index: Nat) -> {Batched.stretch(index, Nat.is_eq(extent, 1n)) == Shape.broadcast_index(extent, index) : Nat}:
348  match extent:
349    case 0n:
350      {==}
351    case 1n+e:
352      match e:
353        case 0n:
354          {==}
355        case 1n+f:
356          {==}