def imap

bendtorch@1.3.0.0 · tensor.bend:155

import Base

155# the position m(j)
156def imap(+m: IMap, +j: U32) -> U32:
157  match m:
158    case Strided{+base, axs}:
159      strided(axs, j, base)
160    case Drop{+inner, +size}:
161      ((U32.div(j, (inner * size : U32)) * inner : U32) + U32.mod(j, inner) : U32)