def dynamic_update_slice(operand, update, start_indices):
  """Updates a dynamic slice.

  See the docstring of `jax.lax.dynamic_update_slice`
  (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dynamic_update_slice.html)
  for details.

  Args:
    operand: an array to slice.
    update: an array containing the new values to write onto `operand`.
    start_indices: a vector of integers, one per dimension. The starts of the
      slice. The vector can be dynamic.

  Returns:
    The updated version of `operand`.
  """
  operand = tf_np.asarray(operand).data
  update = tf_np.asarray(update).data
  start_indices = tf_np.asarray(start_indices, np.int32).data
  if not update.shape.is_fully_defined():
    raise ValueError("update's shape must be fully defined")
  slice_sizes = update.shape
  idx = _get_dynamic_indices(operand, start_indices, slice_sizes)
  if idx is None:
    # `np.zeros([])[()] = 1.0` will result in a scalar array of 1.0
    return tf_np.asarray(update)
  operand = tf.tensor_scatter_nd_update(operand, idx, update)
  return tf_np.asarray(operand)
Updates a dynamic slice.

  See the docstring of `jax.lax.dynamic_update_slice`
  (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.dynamic_update_slice.html)
  for details.

  Args:
    operand: an array to slice.
    update: an array containing the new values to write onto `operand`.
    start_indices: a vector of integers, one per dimension. The starts of the
      slice. The vector can be dynamic.

  Returns:
    The updated version of `operand`.
  
trax/tf_numpy/extensions/extensions.py