[Afloat16]

Problem

ClassWeightedReducer casts stored class weights to the loss dtype before multiplication. That prematurely rounds weights and can zero a small coefficient before its product with a large loss becomes representable. It also computes products such as 40000 * 2 in float16 before averaging, so a representable batch mean can become infinite.

Changes

Convert weights directly to the accumulation dtype, multiply and average float16/bfloat16 losses in float32, and return the scalar in the original loss dtype. The registered weight buffer keeps its precision. Tests compare independent float64 values and gradients, cover reordered anchors and small weights, and retain float16 overflow coverage.

The component documentation, numerical rationale, and unreleased notes describe the behavior.

Validation

• Python 3.12.14 and PyTorch 2.15.0+cpu.
• TEST_DEVICE=cpu TEST_DTYPES=float32,float64,bfloat16 OMP_NUM_THREADS=1 python -m pytest tests/reducers/test_class_weighted_reducer.py -q: 10 passed.
• Ruff lint/format checks and git diff --check passed.

The checksum-verified pinned Code Polishy 0.31.10 commands were attempted. check --git-changes is blocked by an unavailable policy-owned CPython 3.12 interpreter/security providers. test --changed --base 0db540f1b8adb0f65ea91075ee1625e13176758b could not start the declared suites because the repository-specific .venv/bin/uv is absent. The direct focused tests below ran in the available Python/PyTorch environment; these results do not constitute a pass of the complete locked verification gates. Float16 execution is also blocked on this host by RuntimeError: Failed to initialize cpuinfo!; its parameterized cases are retained for supported CI hosts.

[Afloat16]

Could a maintainer rerun the cancelled macOS job for current head 65226a7fd5e253a6f1f2e246a0c446263d953a17? The complete verify job, Linux CPU suites on Python 3.12/3.13/3.14, and Windows CPU suite have passed. GitHub shows no executed steps for the cancelled macOS job; this is the remaining incomplete CI lane.