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)