def win_axes

bendtorch@1.3.0.0 · tensor.bend:306

import Base

306# the index map of reading a tensor of strides st (given by its shape st) at every position of a tensor of shape o, from
307# a base: axis e has o's size and st's stride (inner(st, e)); the last axis first
308def win_axes(n: Nat, +o: List<&2, Nat>, +st: List<&2, Nat>, +e: Nat, acc: List<&2, Ax>) -> List<&2, Ax>:
309  match n:
310    case 0n:
311      acc
312    case 1n+p:
313      win_axes(p, o, st, 1n+e, Ax{U32.from_nat(at1(o, e)), U32.from_nat(inner(st, e))} <> acc)