def modified_cantor_pairing(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
  """Given integer tensors a and b, produces tensor of hashes."""
  # https://en.wikipedia.org/wiki/Pairing_function#Cantor_pairing_function
  # The function produces a modification of the original function making it
  # work consistently under modular 2^32 and 2^64. Namely, we remove // 2 from
  # (a + b) * (a + b + 1) // 2 + b
  b = tf.cast(b, dtype=a.dtype)
  return (a + b) * (a + b + 1) + b
Given integer tensors a and b, produces tensor of hashes.
language/mentionmemory/utils/mention_preprocess_utils.py