law slice_plan_axis

bendtorch-core@0.1.0.0 · src/tensor/slice_LAWS.bend:140

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice

140# The plan of a request at axis `dim` of prefix ++ [extent] ++ suffix, with
141# `dim` the rank of the prefix, holds the request's own start, stop and step,
142# clamped to that extent, whatever the other extents, an empty input's
143# included.
144law slice_plan_axis:
145  for +prefix: List<&2, U32>
146  for +extent: U32
147  for +suffix: List<&2, U32>
148  for +dim: U32
149  for at: {U32.to_nat(dim) == Shape.rank_dimensions(prefix) : Nat}
150  for +start: U32
151  for +stop: U32
152  for +step: U32
153  {axis_part(Slice.plan_of(List.append(&2, U32, prefix, extent <> suffix), dim, start, stop, step)) ==
154    axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)) : Nat & Nat & Nat & Nat}