Fix masking in label smoothing

[tandede]

Summary

• distribute label-smoothing mass only across categories selected by where
• keep fully masked groups finite by using a nonzero normalization denominator
• add JIT regression coverage for partial and fully masked category sets

Test plan

• pre-commit run --files optax/losses/_smoothing.py optax/losses/_smoothing_test.py
• python -m flake8 on the changed files
• python -m pylint on the changed files (10.00/10)
• pytest -q optax/losses/_smoothing_test.py (3 passed)
• pytest -q -n 4 optax (2617 passed, 78 skipped, 3382 subtests passed)

[tandede]

Hi maintainers, I’m following up on this PR in case it was missed. The CLA, security scans, linting, type checks, doctests, and Python/JAX test matrix have all passed. The label-smoothing masking fix is ready for review, and I’d be happy to address any feedback. Thank you!

[tandede]

I rebased this branch onto the latest main and resolved the conflict. The conflicting Pyrefly fixes had already landed upstream, so I dropped the duplicate helper commit and retained only the label-smoothing fix and its regression test.

I reran the full Optax test suite locally: 2,617 passed, 78 skipped, with 3,382 subtests passed. Pre-commit, Ruff, Pyrefly, and git diff --check also passed. The current head is 9ae27d9.