def reduce_batch_count_l_moments(
    x: common_types.TensorType, reduce_instance_dims: bool
) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor,
           tf.Tensor, tf.Tensor]:
  """Computes element first 4 L-moments and the corresponding counts.

  Computes the first 4 L-moments (https://en.wikipedia.org/wiki/L-moment) and
  the number of samples, pairs, etc. used to compute them.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
        dimensions to arrive at a single scalar output. Otherwise, only
        collapses the batch dimension and outputs a `Tensor` of the same shape
        as the input.

  Returns:
    The tuple (count_samples, l1, count_pairs, l2, count_triplets, l3,
    count_quadruplets, l4). Each entry is a `Tensor` with the same dtype as x.
    If reduce_instance_dims is True, the tensors are scalars; otherwise the
    shape is x.shape[1:], i.e. the batch dimension is removed.
  """
  if isinstance(x, tf.SparseTensor) and reduce_instance_dims:
    x = x.values
  elif isinstance(x, tf.RaggedTensor):
    if reduce_instance_dims:
      x = x.flat_values
    else:
      raise NotImplementedError(
          'L-moments only support reduced dims for RaggedTensors')

  if isinstance(x, tf.SparseTensor):
    batch_size = x.dense_shape[0]
    x_rank_2 = tf.sparse.reshape(x, [batch_size, -1])
    dim_1 = x_rank_2.dense_shape[1]
    initial_values = tf.zeros([dim_1], dtype=x.dtype)
    (unused_current_index, l1_sum, l2_sum, l3_sum, l4_sum,
     count_samples, count_pairs, count_triplets,
     count_quadruplets, unused_x_rank_2) = tf.while_loop(
         _condition_l_moments_sparse,
         _iteration_l_moments_sparse,
         [tf.constant(0, dim_1.dtype)] + [initial_values] * 8 + [x_rank_2])
    if reduce_instance_dims:
      final_shape = ()
    elif x.get_shape().ndims and x.get_shape()[1:].is_fully_defined():
      final_shape = x.get_shape()[1:]
    else:
      final_shape = tf.shape(x)[1:]
    l1 = tf.reshape(l1_sum, final_shape)
    l2 = tf.reshape(l2_sum, final_shape)
    l3 = tf.reshape(l3_sum, final_shape)
    l4 = tf.reshape(l4_sum, final_shape)
    count_l1 = tf.reshape(count_samples, final_shape)
    count_l2 = tf.reshape(count_pairs, final_shape)
    count_l3 = tf.reshape(count_triplets, final_shape)
    count_l4 = tf.reshape(count_quadruplets, final_shape)

  else:
    num_samples = tf.size(x) if reduce_instance_dims else tf.shape(x)[0]
    (count_samples, count_pairs, count_triplets, count_quadruplets,
     l1_factors, l2_factors, l3_factors, l4_factors) = _num_terms_and_factors(
         num_samples, x.dtype)
    x_rank_2 = tf.reshape(x, [num_samples, -1])
    dim_1 = tf.shape(x_rank_2)[1]
    initial_moment_values = tf.zeros([dim_1], dtype=x.dtype)
    (unused_current_index, l1_sum, l2_sum, l3_sum, l4_sum, unused_l1_factors,
     unused_l2_factors, unused_l3_factors, unused_l4_factors,
     unused_x_rank_2) = tf.while_loop(
         _condition_l_moments_dense,
         _iteration_l_moments_dense,
         [tf.constant(0, dim_1.dtype)] + [initial_moment_values] * 4 +
         [l1_factors, l2_factors, l3_factors, l4_factors, x_rank_2])
    final_shape = (() if reduce_instance_dims else tf.shape(x)[1:])
    l1 = tf.reshape(l1_sum, final_shape)
    l2 = tf.reshape(l2_sum, final_shape)
    l3 = tf.reshape(l3_sum, final_shape)
    l4 = tf.reshape(l4_sum, final_shape)
    count_l1 = tf.fill(final_shape, count_samples)
    count_l2 = tf.fill(final_shape, count_pairs)
    count_l3 = tf.fill(final_shape, count_triplets)
    count_l4 = tf.fill(final_shape, count_quadruplets)

  return  (count_l1, l1, count_l2, l2, count_l3, l3, count_l4, l4)
Computes element first 4 L-moments and the corresponding counts.

  Computes the first 4 L-moments (https://en.wikipedia.org/wiki/L-moment) and
  the number of samples, pairs, etc. used to compute them.

  Args:
    x: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A bool, if True - collapses the batch and instance
        dimensions to arrive at a single scalar output. Otherwise, only
        collapses the batch dimension and outputs a `Tensor` of the same shape
        as the input.

  Returns:
    The tuple (count_samples, l1, count_pairs, l2, count_triplets, l3,
    count_quadruplets, l4). Each entry is a `Tensor` with the same dtype as x.
    If reduce_instance_dims is True, the tensors are scalars; otherwise the
    shape is x.shape[1:], i.e. the batch dimension is removed.
  
tensorflow_transform/tf_utils.py