def narrow

bendtorch@1.3.0.0 · tensor.bend:540

import Base

540# positions st .. st + l - 1 of axis d
541def narrow(+s: List<&2, Nat>, +d: Nat, +st: Nat, +l: Nat, -pf: {True{} == narrow_ok(s, d, st, l) : Bool}, x: Tensor<s>) -> TP<s, setd(s, d, l)>:
542  match x:
543    case Tensor{dx}:
544      tp_of(s, setd(s, d, l), gather_l(numel(setd(s, d, l)), 0, narrow_map(s, d, st, l), Array.get(F32, dx, imap(narrow_map(s, d, st, l), 0)), new(numel(setd(s, d, l)))))