def _get_dynamic_indices(operand, start_indices, slice_sizes):
  """Calcuates the indices for `tf.gather_nd` from slices.

  Args:
    operand: a Tensor to slice.
    start_indices: a vector Tensor 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 index array suitable for `tf.gather_nd` and `tf.scatter_nd`, or `None` if
    `operand` is a scalar.
  """
  rank = len(slice_sizes)
  operand_rank = tf.rank(operand)
  tf.debugging.Assert(operand_rank == rank, [operand_rank, rank])
  starts_rank = tf.rank(start_indices)
  tf.debugging.Assert(starts_rank == 1, [starts_rank])
  num_starts = tf.shape(start_indices)[0]
  tf.debugging.Assert(num_starts == rank, [num_starts, rank])
  operand_shape = tf.shape(operand)
  tf.debugging.Assert(tf.reduce_all(slice_sizes <= operand_shape),
                      [slice_sizes, operand_shape])
  if rank == 0:
    return None
  start_indices = tf.where(
      start_indices < 0, start_indices + operand_shape, start_indices)
  idx_list = []
  for i in range(rank):
    start = start_indices[i]
    size = slice_sizes[i]
    dim = operand_shape[i]
    start = tf.clip_by_value(start, 0, dim - size)
    # XLA requires tf.range's `start` to be compile-time constant, so we can't
    # do tf.range(start, ...).
    idx = start + tf.range(size)
    shape = [1] * rank
    shape[i] = size
    idx = tf.reshape(idx, shape)
    idx_list.append(idx)
  slice_sizes_tensor = tf.convert_to_tensor(slice_sizes)
  # tf.stack doesn't support broadcasting, so we need to broadcast manually.
  # TODO(wangpeng): Reduce peak memory by broadcasting one-by-one instead of
  #   all-together.
  idx_list = [tf.broadcast_to(x, slice_sizes_tensor) for x in idx_list]
  return tf.stack(idx_list, axis=-1)
Calcuates the indices for `tf.gather_nd` from slices.

  Args:
    operand: a Tensor to slice.
    start_indices: a vector Tensor 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 index array suitable for `tf.gather_nd` and `tf.scatter_nd`, or `None` if
    `operand` is a scalar.
  
trax/tf_numpy/extensions/extensions.py