def dynamic_slice_in_dim(operand, start_index, slice_size, axis=0):
  """Convenience wrapper around dynamic_slice applying to one dimension."""
  operand = tf_np.asarray(operand)
  start_indices = [0] * operand.ndim
  slice_sizes = list(operand.shape)
  axis = int(axis)
  start_indices[axis] = start_index
  slice_sizes[axis] = int(slice_size)
  return dynamic_slice(operand, start_indices, slice_sizes)
Convenience wrapper around dynamic_slice applying to one dimension.
trax/tf_numpy/extensions/extensions.py