def batch_dims

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

import Base

94# The output's batch dimensions: as many as the operand with more of them.
95def batch_dims(+rank: Nat, +first: U32, +second: U32) -> List<&2, U32>:
96  match rank:
97    case 0n:
98      []
99    case 1n+rest:
100      match rest:
101        case 0n:
102          [second]
103        case 1n+more:
104          [first, second]