def METHOD_NAME(t, axis):
  """Squeeze composite tensor along axis `axis`.
  Args:
    t: A `Tensor` or `SparseTensor`.
    axis: A python integer.
  Returns:
    The tensor with dimension `axis` removed.
  Raises:
    InvalidArgumentError: If `t` is a `SparseTensor` and has more than one index
    stored along `axis`.
  """
  if isinstance(t, tf.SparseTensor):
    # Fill in a dummy value if there are no elements in the tensor.
    indices_axis = t.indices[:, axis]
    all_zero = tf.reduce_all(tf.equal(indices_axis, 0))
    with tf.control_dependencies(
        [
            tf.Assert(
                all_zero,
                [
                    (
                        'Unable to squeeze SparseTensor {} axis {} '
                        'because indices are not all equal to 0:'
                    ),
                    indices_axis,
                ],
            )
        ]
    ):
      return tf.SparseTensor(
          indices=tf.concat(
              (t.indices[:, :axis], t.indices[:, axis + 1 :]), axis=1
          ),
          values=t.values,
          dense_shape=tf.concat(
              (t.dense_shape[:axis], t.dense_shape[axis + 1 :]), axis=0
          ),
      )
  else:
    return tf.METHOD_NAME(t, [axis])