BatchedMatmul.batch_dims: def of bendtorch-core@0.1.0.0
Import: import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as BatchedMatmul
Lines 94-104 of bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend&start_line=94&end_line=104

import Base

# The output's batch dimensions: as many as the operand with more of them.
def batch_dims(+rank: Nat, +first: U32, +second: U32) -> List<&2, U32>:
  match rank:
    case 0n:
      []
    case 1n+rest:
      match rest:
        case 0n:
          [second]
        case 1n+more:
          [first, second]

Uses (3):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32