def _set_dim_0(t, d0):
    """Sets the 0-th dimension of the input tensor.

  Args:
    t: the input tensor, assuming the rank is at least 1.
    d0: an integer indicating the 0-th dimension of the input tensor.

  Returns:
    the tensor t with the 0-th dimension set.
  """
    t_shape = t.get_shape().as_list()
    t_shape[0] = d0
    t.set_shape(t_shape)
    return t
