def dynamic_slice(operand, start_indices, slice_sizes):
  """Slicing operation where the indices can be dynamic vlaues.

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

  Args:
    operand: an array to slice.
    start_indices: a vector of integers, one per dimension. The starts of the
      slice. The vector can be dynamic.
    slice_sizes: a list of integers, one per dimension. The sizes of the slice.

  Returns:
    An array containing the slice, with shape equal to `slice_sizes`.
  """
  # This implementation uses tf.gather_nd to implement dynamic_slice, which is
  # memory inefficient because the size of `indices` given to gather_nd is
  # large.
  operand = tf_np.asarray(operand).data
  start_indices = tf_np.asarray(start_indices, np.int32).data
  idx = _get_dynamic_indices(operand, start_indices, slice_sizes)
  if idx is not None:
    operand = tf.gather_nd(operand, idx)
  return tf_np.asarray(operand)
Slicing operation where the indices can be dynamic vlaues.

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

  Args:
    operand: an array to slice.
    start_indices: a vector of integers, one per dimension. The starts of the
      slice. The vector can be dynamic.
    slice_sizes: a list of integers, one per dimension. The sizes of the slice.

  Returns:
    An array containing the slice, with shape equal to `slice_sizes`.
  
trax/tf_numpy/extensions/extensions.py