fix(attention test): give the mask tests the axis rules the cp gate reads

[gulsumgudukbay]

Description

Three mask tests build a 2- or 4-device context mesh and assert the mask is derived from the load-balanced segment_positions. That path is guarded by _load_balanced_context_parallel(), which sizes cp by resolving
activation_q_length against the ambient flax rules - deliberately, so that eval running under logical_axis_rules_for_eval cannot claim a DUAL_CHUNK_SWAP order the input pipeline never applied.

The tests never enter an axis_rules context, so the logical name resolves to P(None), cp_size comes out as 1, load balancing is skipped, and the op returns a plain index-order causal mask while the tests still expect the position-aware one. This makes the activation_q_length rule ambient around the op, the way the rest of the suite does, and sets ulysses_context_sharding, which the cp-size computation reads once the query axis actually resolves.

In the two LoadBalancedMaskTest cases only that one rule is made ambient rather than the whole stub config: those configs also declare segment_ids_batch -> context, which was inert while nothing made it ambient, and cannot hold here because generate_attention_mask would shard their batch-1 segment ids across the 4-way context mesh.

These tests skip unless 2+ devices are visible, and they are auto-marked cpu_only so the 4-GPU gpu-unit job deselects them - which is why CI has not been executing them at all.

Tests

Reproduce the failure on main and the fix on this branch with:

XLA_FLAGS=--xla_force_host_platform_device_count=4 python -m pytest \
  tests/unit/attention_test.py -k "uses_original_load_balanced_positions or uses_segment_positions"

Verified on 4x AMD Instinct MI355X (ROCm 10, jax 0.11.1): the three tests fail on main (plain causal mask,
then an IndivisibleError once the rule resolves) and pass on this branch. The full tests/unit/attention_test.py
file also passes there (72 passed, 82 skipped, 11 subtests passed).

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Code Review

This pull request updates several unit tests in tests/unit/attention_test.py to support Ulysses context sharding and logical axis rules. Specifically, it adds ulysses_context_sharding and activation_q_length axis rules to the test configurations, and wraps the AttentionOp initialization and mask generation within the nn_partitioning.axis_rules context manager. I have no feedback to provide as there are no review comments to evaluate.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!