jax.lax.dynamic_update_index_in_dim#

jax.lax.dynamic_update_index_in_dim(operand, update, index, axis)[source]#

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

Parameters
Return type

Array