def kept_from

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

import Base

37# The dimensions from axis `axis` on, each reduced one set to 1.
38def kept_from(+axes: List<&2, U32>, dimensions: List<&2, U32>, +axis: U32) -> List<&2, U32>:
39  match dimensions:
40    case Nil{}:
41      Nil{}
42    case extent <> rest:
43      Bool.pick(U32, reduces(axes, axis), 1, extent) <> kept_from(axes, rest, U32.add(axis, 1))