def picked

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

import Base

246# The gradient index input cell `position` reads, or None where no output
247# cell came from it: position is (o * extent + i) * inner + r, and it reads
248# (o * length + k) * inner + r when i = start + k * step with k below the
249# length, for k = (i - start) / step.
250def picked(+plan: Plan, +row: Nat, +within: Nat, +i: Nat, +k: Nat, hit: Bool) -> Maybe<&2, Nat>:
251  Plan{+outer, +extent, +inner, +start, +step, +length} = plan
252  match hit:
253    case False{}:
254      None{}
255    case True{}:
256      Some{Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(row, extent), length), k), inner), within)}