Generate a mask to prevent the batch to attend to each others.

  Args:
    batch_coordinates_q: Int-like Tensor of shape [length_q, 1] containing the
      coordinates of the batches
    batch_coordinates_k: Int-like Tensor of shape [length_k, 1] containing the
      coordinates of the batches. If None, do self-attention.
    condition_fn: Callable defining the attention mask.

  Returns:
    Float-like Tensor of shape [length_q, length_k] containing either 0 or
    -infinity (-1e9).

def attention_bias_batch(batch_coordinates_q,
                         batch_coordinates_k=None,
                         condition_fn=None):
  """Generate a mask to prevent the batch to attend to each others.

  Args:
    batch_coordinates_q: Int-like Tensor of shape [length_q, 1] containing the
      coordinates of the batches
    batch_coordinates_k: Int-like Tensor of shape [length_k, 1] containing the
      coordinates of the batches. If None, do self-attention.
    condition_fn: Callable defining the attention mask.

  Returns:
    Float-like Tensor of shape [length_q, length_k] containing either 0 or
    -infinity (-1e9).
  """
  if batch_coordinates_k is None:
    batch_coordinates_k = batch_coordinates_q

  # Convert to float first because of b/25387198.
  def to_float(bc):
    bc = tf.squeeze(bc, 1)
    bc = tf.to_float(bc)
    return bc

  # Broadcast to create [length_q, length_k] mask.
  bc_v = tf.expand_dims(to_float(batch_coordinates_q), 1)
  bc_h = tf.expand_dims(to_float(batch_coordinates_k), 0)
  bias_batch = bc_h - bc_v
  bias_batch = condition_fn(bias_batch)
  bias_batch *= -1e9
  return bias_batch