jax.lax.dynamic_slice

jax.lax.dynamic_slice(operand, start_indices, slice_sizes)[source]

Wraps XLA’s DynamicSlice operator.

Parameters
  • operand (Any) – an array to slice.

  • start_indices (Sequence[Any]) – a list of scalar indices, one per dimension. These values may be dynamic.

  • slice_sizes (Sequence[int]) – the size of the slice. Must be a sequence of non-negative integers with length equal to ndim(operand). Inside a JIT compiled function, only static values are supported (all JAX arrays inside JIT must have statically known size).

Return type

Any

Returns

An array containing the slice.