def _to_term_document_one_hot(
    x: tf.SparseTensor, vocab_size: Union[int, tf.Tensor]) -> tf.SparseTensor:
  """Creates a one-hot SparseTensor of term existence for every doc/term pair.

  Converts a <batch_indices, local_sequence_index>-indexed, <vocab_index>-valued
  sparse tensor to <batch_indices, vocab_index> one-hot tensor to represent the
  existence of each vocab term in each document of a batch. For example, when x
  has the dense form:
    [[3, 2, 3],  # first example of the batch has vocab term 2 and 3
     [1, 1]],    # second example of the batch has vocab term 1
  with vocab_size=4, the dense form of the out one-hot tensor is
    [[0, 0, 1, 1],
     [0, 1, 0, 0]]

  Args:
    x: a SparseTensor of int64 representing string indices in vocab. The indices
      are <batch_index, local_index> and the values are <vocab_index>.
      Typically, x is the output of tft.compute_and_apply_vocabulary.
    vocab_size: A scalar int64 Tensor - the count of vocab used to turn the
      string into int64s including any OOV buckets.

  Returns:
    a SparseTensor with size (batch_size, vocab_size), indices being
      <doc_index_in_batch>, <term_index_in_vocab> and int32 values being 1 for
      all mentioned terms or 0 if not shown in each document.
  """
  vocab_size = tf.convert_to_tensor(value=vocab_size, dtype=tf.int64)

  # Combine batch indices (first column of x's indices) and vocab indices (
  # x's values) as new indices (<batch_index>, <vocab_index>).
  batch_indices = x.indices[:, 0]  # sparse tensor indices are int64
  vocab_indices = tf.cast(x.values, dtype=tf.int64)

  # Dedup (<batch_index>, <vocab_index>) pairs. This is because document
  # frequency only cares the existence of a term in a document, not the
  # occurrence frequency within that document.
  # Hashing (<batch_index>, <vocab_index>) pairs for dedup.
  multiplier = vocab_size + 1
  unique_flatten_indices, _ = tf.raw_ops.UniqueV2(
      x=batch_indices * multiplier + vocab_indices, axis=[0])
  unique_batch_indices = tf.cast(
      tf.math.divide(unique_flatten_indices, multiplier), dtype=tf.int64)
  unique_vocab_indices = tf.math.mod(unique_flatten_indices, multiplier)
  unique_batch_vocab_indices = tf.transpose(
      tf.stack([unique_batch_indices, unique_vocab_indices]))

  # If term i shows at least once in document j, then doc_freq<i, j> = 1
  one_hot_values = tf.ones_like(unique_flatten_indices, dtype=tf.int32)

  # New shape of the one hot tensor is batch_size * vocab_size
  new_shape = tf.stack([x.dense_shape[0], vocab_size])
  new_shape.set_shape([2])

  return tf.SparseTensor(
      indices=unique_batch_vocab_indices,
      values=one_hot_values,
      dense_shape=new_shape)
Creates a one-hot SparseTensor of term existence for every doc/term pair.

  Converts a <batch_indices, local_sequence_index>-indexed, <vocab_index>-valued
  sparse tensor to <batch_indices, vocab_index> one-hot tensor to represent the
  existence of each vocab term in each document of a batch. For example, when x
  has the dense form:
    [[3, 2, 3],  # first example of the batch has vocab term 2 and 3
     [1, 1]],    # second example of the batch has vocab term 1
  with vocab_size=4, the dense form of the out one-hot tensor is
    [[0, 0, 1, 1],
     [0, 1, 0, 0]]

  Args:
    x: a SparseTensor of int64 representing string indices in vocab. The indices
      are <batch_index, local_index> and the values are <vocab_index>.
      Typically, x is the output of tft.compute_and_apply_vocabulary.
    vocab_size: A scalar int64 Tensor - the count of vocab used to turn the
      string into int64s including any OOV buckets.

  Returns:
    a SparseTensor with size (batch_size, vocab_size), indices being
      <doc_index_in_batch>, <term_index_in_vocab> and int32 values being 1 for
      all mentioned terms or 0 if not shown in each document.
  
tensorflow_transform/experimental/mappers.py