Fix vocab tiling hidden-state cotangent scaling

[JKSPL]

Description

Fix the Linen and NNX vocabulary-tiling custom VJPs so the incoming loss
cotangent scales the hidden-state gradient as well as the output-head gradient.

Each tiled backward pass computes per-tile VJPs with a scalar seed of 1.0 and
then applies the incoming loss_cotangent to the accumulated output-head
gradient. Previously, the same chain-rule factor was omitted from the emitted
hidden-state gradient. Directly differentiating the tiled loss supplies an
implicit cotangent of one, so the existing parity tests could not observe the
error. When a caller normalizes or weights the loss, output-head gradients were
scaled correctly while gradients flowing through the transformer body were
not.

Apply the missing factor in both implementations. Add focused Linen and NNX
regressions that differentiate each tiled loss beneath outer scales of 1.0
and 2.0, then require the hidden-state gradient to double. These mathematical
checks now run in the CPU suite; they do not require accelerator hardware.

Tests

• uv run --no-sync pytest -q tests/unit/tiling_test.py — 2 passed, 18 skipped
• uv run --no-sync pre-commit run pyink --files src/maxtext/utils/vocabulary_tiling.py tests/unit/tiling_test.py — passed
• uv run --no-sync pre-commit run codespell --files src/maxtext/utils/vocabulary_tiling.py tests/unit/tiling_test.py — passed
• uv pip check --python .venv/bin/python — passed
• Red/green verification: both focused regression tests fail when the two new
  cotangent multiplications are removed and pass when restored.

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 and provided workload links above if applicable. End-to-end testing is not applicable to this focused chain-rule fix; numerical unit tests cover both implementations.
• No documentation change is needed for this mathematical bug fix.

[gemini-code-assist]

Code Review

This pull request ensures that the hidden-state gradients in vocabulary tiling correctly scale with the outer loss scale by multiplying grad_reshaped_hidden_states by loss_cotangent in the backward passes. It also adds corresponding unit tests for both Linen and NNX implementations to verify this behavior. The review feedback suggests replacing the in-place multiplication operator (*=) with explicit assignment to adhere to JAX's immutable array paradigm and ensure consistency.

[NuojCheng]

Could you update the pyink in tests/unit/attention_test.py? see https://github.com/AI-Hypercomputer/maxtext/actions/runs/32472985540/job/96840648989?pr=4959

[JKSPL]

Could you update the pyink in tests/unit/attention_test.py? see https://github.com/AI-Hypercomputer/maxtext/actions/runs/32472985540/job/96840648989?pr=4959

done

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!