def pick_at

bendtorch-core@0.1.0.0 · src/tensor/permute_build.bend:37

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

37# The pick of coordinate k of `Shape.digits(dims, x)`; None past the last
38# dimension.
39def pick_at(+dims: List<&2, Nat>, k: Nat) -> Maybe<&2, Pick>:
40  match dims:
41    case Nil{}:
42      None{}
43    case head <> tail:
44      match k:
45        case 0n:
46          Some{Pick{Nil{}, Shape.numel(tail)}}
47        case 1n+j:
48          widen(Shape.numel(tail), pick_at(tail, j))