def batch_nats

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

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

643# The plan's batch extents as naturals, as `Batched.batch_dims` lists them.
644def batch_nats(+rank: Nat, +first: Nat, +second: Nat) -> List<&2, Nat>:
645  match rank:
646    case 0n:
647      []
648    case 1n+rest:
649      match rest:
650        case 0n:
651          [second]
652        case 1n+more:
653          [first, second]