def limit

bendtorch-core@0.1.0.0 · src/autograd/sum_dims.bend:66

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

66# How many axes an input of these dimensions has for PyTorch's check: its
67# rank, and 1 for a rank-0 input.
68def limit(dimensions: List<&2, U32>) -> U32:
69  match dimensions:
70    case Nil{}:
71      1
72    case +extent <> +rest:
73      U32.from_nat(Shape.rank_dimensions(extent <> rest))