jax.lax.dynamic_update_index_in_dim
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 singleaxis
.
Convenience wrapper around dynamic_update_slice()
to update a slice
of size 1 in a single axis
.