chore: scope the jaxtyping lint ignores so F821 runs everywhere else

[mikasenghaas]

Summary

F821 (undefined name) was disabled repo-wide:

select = ["F", "I"]
ignore = ["F722", "F821"] # Need to ignore for jaxtyping

jaxtyping needs it — Int[Tensor, "batch seq"] is a string annotation ruff reads as a forward
reference — but the ignore was global, so ruff could not report a missing import anywhere in
prime-rl, and style.yaml runs only ruff. A missing import is then invisible to lint, invisible
to CI, and invisible to pytest unless a test happens to execute that line: importing the module
succeeds, because the NameError only fires when the function is called.

Four modules actually use jaxtyping annotations, so the ignore names them:

[tool.ruff.lint.per-file-ignores]
"src/prime_rl/trainer/model.py" = ["F722", "F821"]
"src/prime_rl/trainer/rl/data.py" = ["F722", "F821"]
"src/prime_rl/trainer/rl/loss.py" = ["F722", "F821"]
"src/prime_rl/trainer/sft/data.py" = ["F722", "F821"]

New jaxtyping usage outside these fails lint until it is added here, which is the point — the entry
should be deliberate.

The fifth file was a real bug

F821 also fired on trainer/models/layers/rotary_emb.py, and not for jaxtyping:

def _compute_default_rope_parameters(config: Optional["PreTrainedConfig"] = None, ...)

Nothing imports that name, so the annotation refers to nothing —
typing.get_type_hints(_compute_default_rope_parameters) raises NameError on main. Fixed by
importing it (transformers is already a hard import in that module), and the hint now resolves.
It is not in the ignore list.

Verification

• ruff check src tests packages — clean.
• The rule catches the real case. A probe file with a missing import:

  with this PR:              F821 Undefined name `EXCLUDE_FIELDS`
  with the global ignore:    All checks passed!
  

  That is the exact shape of a bug that reached a live run on feat(orchestrator)!: align multi-agent types #3183 through a green suite.
• typing.get_type_hints on the rope helper resolves after the fix, raises before it.
• uv run pytest tests/unit — 485 passed. One pre-existing, unrelated failure on this box
  (test_qwen3_vl_e2e, stale fixture, fix: token_id-formatted logprob tokens in the qwen3-vl fake engine #3161).

Follow-up

This closes the cheap half. The other half is that prime-rl has no type checker in CI at all —
verifiers runs Ty, and something equivalent here would catch the attribute-level version of this
(reading a field that no longer exists), which F821 does not.

🤖 Generated with Claude Code

---

Note

Low Risk
Lint configuration and a type-hint import fix; no runtime behavior change beyond correct annotation resolution.

Overview
Ruff F722/F821 are no longer ignored repo-wide. They stay suppressed only on the four trainer modules that use jaxtyping shape annotations, so missing imports and undefined names surface everywhere else in CI.

rotary_emb.py had a real F821: the helper annotated Optional["PreTrainedConfig"] without importing that name. The PR imports PretrainedConfig from transformers and uses it in the hint so get_type_hints resolves correctly.

^{Reviewed by Cursor Bugbot for commit 062e837. Bugbot is set up for automated code reviews on this repo. Configure here.}