def all_compare_without_pad(xs: Array, ys: Array) -> Array:
  """Performs all-to-all comparison of two 1D arrays ignoring pad values."""
  xs_expanded = jnp.expand_dims(xs, 1)
  ys_expanded = jnp.expand_dims(ys, 0)
  same_ids = xs_expanded == ys_expanded
  same_ids = jnp.logical_and(same_ids, xs_expanded != 0)
  same_ids = jnp.logical_and(same_ids, ys_expanded != 0)
  return same_ids
Performs all-to-all comparison of two 1D arrays ignoring pad values.
language/mentionmemory/utils/mention_utils.py