Fix pre-commit (ruff) and mypy errors

[runame]

Summary

Resolves the remaining pre-commit and type-check job failures on main after PRs #256 (init imports) and #257 (examples workflow). Three commits, each separately reviewable:

Commit 1: lint fixes (ruff E402 + E731) and copyright header normalization

• E402 (~92 violations across 12 files): every affected file had a duplicate triple-quoted copyright block at the top followed by the real module docstring (or, in gpa/, a #-comment copyright). Python only treats the first triple-quoted string as a docstring, so the duplicate counted as a non-import statement and triggered E402 for every subsequent import. Drop the duplicate so the real module docstring is the first statement.
• E731 (2 violations) in _shampoo_fully_shard_lossless_distributor.py:69 and _shampoo_hybrid_shard_lossless_distributor.py:75: rewrite should_assign_param_idx = lambda i: ... as a def.
• Copyright convention: 79/79 untouched files use a single triple-quoted block. The 6 gpa files I touched (plus gpa/gpa_adamw.py, which had both styles) are normalized to that convention by dropping their #-comment copyright in favor of the triple-quoted form.

Commit 2: ruff-format pass

The pinned ruff-format (v0.8.0 in .pre-commit-config.yaml) considers 14 unrelated files mis-formatted (purely whitespace around assert ..., (...) wrapping). Applying the formatter so the hook passes on the first run instead of failing after auto-modifying files.

Commit 3: mypy fixes (33 errors → 0)

All 27 errors reported by CI plus 6 local-only errors that newer mypy/torch stubs surface:

• gpa/gpa_adamw.py (10 errors):
  • step(): add @overload pair to match torch.optim.Optimizer
  • step(): annotate per-group buffer lists as list[Tensor]
  • step(): skip empty parameter groups via continue; rename group-local first_param to group_first_param to avoid mypy union with the Optional[Parameter] from the train-mode check loop
• gpa/tests/gpa_test_utils.py:42: annotate devices as tuple[device, ...]
• gpa/gpu_tests/gpa_adamw_numerics_test.py:31: ignore missing parameterized stubs (no public stub package)
• gpa/examples/cifar10_example.py:115: ignore len() on Dataset[Any]
• distributed_shampoo/distributed_shampoo.py:
  • lines 669-671: drop redundant re-annotations on the AdaGradPreconditionerConfig branch (no-redef)
  • line 1632: move stray # type: ignore onto the _pre_load_train_modes assignment line so it actually applies
• distributed_shampoo/distributor/_shampoo_hybrid_shard_lossless_distributor.py:160: replace filter(lambda, ...) with a generator expression to dodge a confusing TypeGuard overload mypy picks for filter()
• distributed_shampoo/distributor/shampoo_{ddp,hsdp,hybrid_shard}_distributor.py: ignore attr-defined for the private DeviceMesh._get_all_submeshes — local-only, surfaces with current torch stubs
• distributed_shampoo/distributor/shampoo_distributor.py:228: ignore narrowing of map(partial(tuple), ...) assignment — local-only
• distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py:376,411 and examples/parallelism.py:177,206: ignore arg-type for partial[FSDPModule] / FSDPModule passed to a Module parameter (FSDPModule mixes in but mypy can't see it)
• distributed_shampoo/examples/utils.py:161: annotate sampler
• distributed_shampoo/examples/tests/utils_test.py:228,240: replace assertIsNotNone with assert ... is not None so mypy narrows _fmt before assertIn
• distributed_shampoo/utils/load_balancing_utils.py:57: extend existing # type: ignore[misc] to also cover call-overload from max(float, floating[Any])

Local verification

$ ruff check .
All checks passed!
$ ruff format --check .
93 files already formatted
$ mypy .
Success: no issues found in 93 source files

Test plan

• CI pre-commit job passes
• CI type-check job passes
• Other CI jobs (tests, gpu-tests, examples) unchanged

Notes

Stacked on top of #256 (fix/init-stale-rename-imports) and #257 (fix/examples-beta2-config). Until those merge, the diff for this PR will appear to include their commits. Merge order: #256 → #257 → this.

🤖 Generated with Claude Code

[meta-codesync]

@hjmshi has imported this pull request. If you are a Meta employee, you can view this in D103082000.

[wz337]

Review automatically exported from Phabricator review in Meta.