jax.lax.dynamic_update_index_in_dim

jax.lax.dynamic_update_index_in_dim(operand, update, index, axis)

Convenience wrapper around dynamic_update_slice() to update a slice of size 1 in a single axis.

Parameters

• operand (Array) –

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

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

• axis (int) –

Return type

Array