def _preprocess_tensors_for_cooccurences(
    x_input, y_input, weights_input, filter_regex
):
  """Wrangles the tensors to make them compatible with AMI computation.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.

  Returns:
    A tuple containing the re-shaped, verified (x, y, weights).
  """
  tf.compat.v1.assert_type(y_input, tf.int64)
  # TODO(b/134075780): Revisit expected weights shape when input is sparse.
  if isinstance(x_input, tf.SparseTensor):
    batch_indices = x_input.indices[:, 0]
    # y and densified x should have the same batch dimension.
    assert_eq = tf.compat.v1.assert_equal(
        tf.shape(y_input)[0], tf.cast(x_input.dense_shape[0], tf.int32))
    with tf.control_dependencies([assert_eq]):
      y = tf.gather(y_input, batch_indices)
    x = x_input.values
  elif isinstance(x_input, tf.RaggedTensor):
    # Each batch instance in x corresponds to a single value in y.
    x_row_indices = _get_ragged_batch_value_rowids(x_input)
    assert_compatible = tf.debugging.assert_greater_equal(
        tf.shape(y_input, out_type=tf.int64)[0], x_input.bounding_shape(axis=0))
    with tf.control_dependencies([assert_compatible]):
      x = tf.ensure_shape(x_input.flat_values, [None])
      y = tf.gather(y_input, x_row_indices)
  else:
    y = y_input
    x = x_input
  if weights_input is None:
    weights = tf.ones_like(x, dtype=tf.float32)
  else:
    x, weights_input = assert_same_shape(x, weights_input)
    weights = weights_input
  y = _broadcast_to_x_shape(x, y)
  x = tf.reshape(x, [-1])
  filter_fn = _make_regex_filter_fn(x, filter_regex)
  x = filter_fn(x)
  y = filter_fn(tf.reshape(y, [-1]))
  weights = filter_fn(tf.reshape(weights, [-1]))
  return x, y, weights
Wrangles the tensors to make them compatible with AMI computation.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.

  Returns:
    A tuple containing the re-shaped, verified (x, y, weights).
  
tensorflow_transform/tf_utils.py