def _slice(tensor, dim, start, end):
  """Slices the tensor along given dimension."""
  # Performs a slice along the dimension dim. E.g. for tensor t of rank 3,
  # _slice(t, 1, 3, 5) is same as t[:, 3:5].
  # For a slice unbounded to the right, set end=0: _slice(t, 1, -3, 0) is same
  # as t[:, -3:].
  rank = tensor.shape.rank
  slices = rank * [slice(None)]
  if end == 0:
    end = None
  slices[dim] = slice(start, end)
  return tensor[slices]