def axis_extent

bendtorch-core@0.1.0.0 · src/tensor/slice.bend:89

import Base

89# The extent of axis `dim`, 0 past the rank. It reads that one dimension and
90# forms no product.
91def axis_extent(dimensions: List<&2, U32>, +dim: Nat) -> Nat:
92  match dimensions dim:
93    case extent <> rest 0n:
94      U32.to_nat(extent)
95    case extent <> rest 1n+next:
96      axis_extent(rest, next)
97    case Nil{} _:
98      0n