Incorrect layout conversion of `ref_space_uid` in atom cross-attention

Tested on c0f97eda2f1f482fd94d3a38bece18c7069b4a5c.

In src/alphafold3/model/network/atom_cross_attention.py, the following code converts batch.ref_structure.ref_space_uid using queries_to_keys:

  keys_ref_space_uid = atom_layout.convert(
      batch.atom_cross_att.queries_to_keys,
      batch.ref_structure.ref_space_uid,
      layout_axes=(-2, -1),
  )

  offsets_valid = (
      queries_ref_space_uid[:, :, None] == keys_ref_space_uid[:, None, :]
  )

However, batch.ref_structure.ref_space_uid is in token layout here. This causes offsets_valid to be incorrect and masks out almost all positions. I tested it with a small input with 1171 atoms, and only 45 rows in offsets_valid contained nonzero values. It looks like the intended version might be to use queries_ref_space_uid as the source tensor.

Could this be confirmed?