code: def attention_mask_same_segment(
      query_segment, memory_segment=None, dtype=tf.float32):
    
    memory_segment = rename_length_to_memory_length(
        memory_segment or query_segment)
    return mtf.cast(mtf.not_equal(query_segment, memory_segment), dtype) * -1e9
docs: Bias for attention where attention between segments is disallowed.
  
    Args:
      query_segment: a mtf.Tensor with shape [..., length_dim]
      memory_segment: a mtf.Tensor with shape [..., memory_length_dim]
      dtype: a tf.dtype
  
    Returns:
      a mtf.Tensor with shape [..., length_dim, memory_length_dim]