jax.lax.slice_in_dim#

jax.lax.slice_in_dim(operand, start_index, limit_index, stride=1, axis=0)[source]#

Convenience wrapper around slice applying to only one dimension.

Parameters

  • operand (Array) –

  • start_index (Optional[int]) –

  • limit_index (Optional[int]) –

  • stride (int) –

  • axis (int) –

Return type

Array