jax.lax.dynamic_slice_in_dim

jax.lax.dynamic_slice_in_dim(operand, start_index, slice_size, axis=0)

Convenience wrapper around dynamic_slice applying to one dimension.

Parameters

• operand (Array) –

• start_index (Union[Array, ndarray, bool_, number, bool, int, float, complex]) –

• slice_size (int) –

• axis (int) –

Return type

Array