def wrap

bendtorch-core@0.1.0.0 · src/tensor/signed.bend:30

import Base

30# PyTorch's wrap of a position among `limit` axes or elements (maybe_wrap_dim):
31# a position d in -limit..limit-1 becomes d + limit when d is negative and d
32# otherwise. A position outside that range stays outside it, at or past
33# `limit`, so the operation refuses it with its own check, code and order:
34# First{k} stays k, and an out-of-range Last{k} becomes `limit`.
35def wrap(+limit: U32, position: Position) -> U32:
36  match position:
37    case First{k}:
38      k
39    case Last{+k}:
40      Bool.pick(U32, U32.is_lt(k, limit), U32.sub(U32.sub(limit, 1), k), limit)