law select_slices

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

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

293# PyTorch's select(dim, index) is the slice index:index + 1 with step 1,
294# reshaped to the input's dimensions without the axis, once the axis is in
295# range and the index below the extent; an index at or past the extent is
296# refused with OutOfRange, an axis past the rank as a slice refuses it.
297law select_slices:
298  for input: Core.Dynamic
299  for +dim: U32
300  for +index: U32
301  for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool}
302  for below: {Nat.is_lt(U32.to_nat(index), extent_of(input, dim)) == True{} : Bool}
303  {Slice.select(input, dim, index) ==
304    reshaped(Slice.dropped(Shape.dimensions(Totality.dynamic_shape(input)), U32.to_nat(dim)), Slice.slice(input, dim, index, U32.add(index, 1), 1)) :
305    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}