def sparse_tensor_left_align(sparse_tensor: tf.SparseTensor) -> tf.SparseTensor:
  """Re-arranges a `tf.SparseTensor` and returns a left-aligned version of it.

  This mapper can be useful when returning a sparse tensor that may not be
  left-aligned from a preprocessing_fn.

  Args:
    sparse_tensor: A 2D `tf.SparseTensor`.

  Raises:
    ValueError if `sparse_tensor` is not 2D.

  Returns:
    A left-aligned version of sparse_tensor as a `tf.SparseTensor`.
  """
  if sparse_tensor.get_shape().ndims != 2:
    raise ValueError('sparse_tensor_left_align requires a 2D input')
  reordered_tensor = tf.sparse.reorder(sparse_tensor)
  transposed_indices = tf.transpose(reordered_tensor.indices)
  row_indices = transposed_indices[0]
  row_counts = tf.unique_with_counts(row_indices, out_idx=tf.int64).count
  column_indices = tf.ragged.range(row_counts).flat_values
  return tf.SparseTensor(
      indices=tf.transpose(tf.stack([row_indices, column_indices])),
      values=reordered_tensor.values,
      dense_shape=reordered_tensor.dense_shape)
Re-arranges a `tf.SparseTensor` and returns a left-aligned version of it.

  This mapper can be useful when returning a sparse tensor that may not be
  left-aligned from a preprocessing_fn.

  Args:
    sparse_tensor: A 2D `tf.SparseTensor`.

  Raises:
    ValueError if `sparse_tensor` is not 2D.

  Returns:
    A left-aligned version of sparse_tensor as a `tf.SparseTensor`.
  
tensorflow_transform/mappers.py