code: def attention_bias_batch(batch_coordinates_q,
                           batch_coordinates_k=None,
                           condition_fn=None):
    
    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
docs: 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).