Preventing leak in packed sequences

[saareliad]

When packing is done here https://github.com/tensorflow/mesh/blob/6a812c8bb847e081e976533ed497c7c5016bb1ec/mesh_tensorflow/transformer/dataset.py
Each packed sequence has multiple examples ("segments"). I'm trying to figure out where do you prevent information to leak between these examples (e.g in attention).

I came across this

mesh/mesh_tensorflow/layers.py

Line 1813 in 4db643b

def attention_mask_same_segment(

But I see it is not used anywhere.

I can't seem to find where the information leak is prevented elsewhere. Can you clarify?