Fix several latent bugs in MSA, ranking, and structure-context code

[mooreneural]

Fixes 5 small but real bugs found while reviewing the data pipeline, ranking, and feature-generator code. Each is a one- to four-line change.

Fixes

1. colabfold.py — error_count = 0 was inside the while True retry loop in submit() and status(), so the counter reset every iteration and if error_count > 5: raise was unreachable. The submit/status helpers would retry forever on persistent server errors. Moved the initialization above the loop. (download() was already correct.)

2. all_atom_structure_context.py — In AllAtomStructureContext.merge, token_backbone_frame_index (which stores atom indices, see backbone_atoms_indices in data/dataset/structure/utils.py) was being offset by token_offsets rather than atom_offsets when merging multiple chains. The sibling fields token_centre_atom_index and token_ref_atom_index in the same function correctly use atom_offsets. Currently latent because the returned frame_idces from get_frames_and_mask is discarded by chai1.py, but it's still a real semantic bug.

3. token_dist_restraint.py — Stray extra [ ... ] around the key_entity_types tensor literal gave it shape (1, n) instead of (n,), asymmetric with query_entity_types defined right above it. Currently unused at the call site, but trivially wrong.

4. ranking/clashes.py — In has_inter_chain_clashes, the two ratio checks used inconsistent einops patterns: "... c -> ... c 1" (ellipsis-form) on one line and "b c -> b 1 c" (hardcoded b) on the next. Functionally equivalent for the 2D batch input used today, but the latter will fail if batch dims ever change. Made both ellipsis-form.

5. features/generators/msa.py — MSADataSourceGenerator._generate mutated the input msa_sequence_source tensor in place (msa_sequence_source[msa_sequence_source.eq(query)] = none) before reassigning it via the non-inplace masked_fill on the next line. Idempotent so it doesn't cause incorrect output today, but fragile — replaced with masked_fill for consistency.

Test plan

• python -m py_compile on all modified files (passes locally)
• Existing tests/test_inference_dataset.py::test_protein_with_smiles exercises the multi-chain merge() code path and still passes
• Manual verification of einops pattern equivalence on 2D input + correctness on 3D input
• Manual verification that key_entity_types now has shape (n,) matching query_entity_types