def axes

bendtorch@1.3.0.0 · tensor.bend:136

import Base

136# the axes of reading a tensor of shape a (aligned to o, same rank) at every position of o; the last axis first
137def axes(ra: List<&2, Nat>, ro: List<&2, Nat>, +st: U32, acc: List<&2, Ax>) -> List<&2, Ax>:
138  match ra ro:
139    case Con{+d, ta} Con{+n, to}:
140      axes(ta, to, (st * U32.from_nat(d) : U32), Ax{U32.from_nat(n), stride_of(d, st)} <> acc)
141    case _ _:
142      acc