def dynamic_update_slice_in_dim(operand, update, start_index, axis):
  """Convenience wrapper around dynamic_update_slice for one dimension."""
  operand = tf_np.asarray(operand)
  axis = int(axis)
  start_indices = [0] * operand.ndim
  start_indices[axis] = start_index
  return dynamic_update_slice(operand, update, start_indices)
Convenience wrapper around dynamic_update_slice for one dimension.
trax/tf_numpy/extensions/extensions.py