def size(x: common_types.TensorType,
         reduce_instance_dims: bool = True,
         name: Optional[str] = None) -> tf.Tensor:
  """Computes the total size of instances in a `Tensor` over the whole dataset.

  Args:
    x: A `Tensor`, `SparseTensor`, or `RaggedTensor`.
    reduce_instance_dims: By default collapses the batch and instance dimensions
      to arrive at a single scalar output. If False, only collapses the batch
      dimension and outputs a vector of the same shape as the input.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor` of type int64.
  """
  with tf.compat.v1.name_scope(name, 'size'):
    # Note: Calling `sum` defined in this module, not the builtin.
    if isinstance(x, tf.SparseTensor):
      ones_like_x = tf.SparseTensor(
          indices=x.indices,
          values=tf.ones_like(x.values, tf.int64),
          dense_shape=x.dense_shape)
    else:
      ones_like_x = tf.ones_like(x, dtype=tf.int64)
    return sum(ones_like_x, reduce_instance_dims)
Computes the total size of instances in a `Tensor` over the whole dataset.

  Args:
    x: A `Tensor`, `SparseTensor`, or `RaggedTensor`.
    reduce_instance_dims: By default collapses the batch and instance dimensions
      to arrive at a single scalar output. If False, only collapses the batch
      dimension and outputs a vector of the same shape as the input.
    name: (Optional) A name for this operation.

  Returns:
    A `Tensor` of type int64.
  
tensorflow_transform/analyzers.py