[qflen]

Fixes #192812

Root cause

exponential_ samples in a wider type and narrows with static_cast<scalar_t>. Anything below float16's smallest subnormal becomes an exact 0.

On CPU the kernel samples in double, so a sample under 2^-25 rounds to zero about 3 times per 1e8 draws at lambda=1. On CUDA the eps/2 floor inside transformation::exponential is applied in accscalar_t (float), which bounds the sample at 2^-24/lambda; that sits at or below half's round-to-even tie point once lambda >= 2, so it narrows to zero as well.

multinomial without replacement divides p by the sample (Gumbel trick), so a zero turns a zero-probability category into 0/0 = NaN, and NaN wins argmax. Only float16 is affected; bfloat16 and float32 subnormals are far below anything either kernel emits.

Fix

c10::cast_no_underflow<To>(value) narrows, and returns denorm_min with the input's sign when the cast would otherwise flush a nonzero value to zero. Both kernels use it in place of their static_cast.

It lives in c10/util/safe_numerics.h rather than TypeCast.h because the template itself needs only <limits> (the destination's numeric_limits is instantiated at the call site), so the header stays cheap enough to include from CUDA kernels.

The random stream is untouched and an exact 0 stays 0, so exponential_(inf) still returns 0 (asserted by test_exponential) and every other output value is bit-identical.

Flooring before the cast does not work, since eps/2 of a double still rounds to float16 0. Sampling on CPU in float with the CUDA floor would change the random stream for fp16/bf16/fp32 and still emit zeros for lambda >= 2, so it would need this floor anyway.

Verification

test_exponential_no_zero is now parametrized over lambda and also runs on CPU. At lambda=1024 roughly 3e-5 of the samples fall under the rounding threshold, so 5e7 samples produce about 1500 zeros before the fix against 1 at lambda=1; that is what makes it a real regression detector on both backends. test_exponential_half_underflow pins a deterministic CPU case (seed 20, index 5833).

python test/test_torch.py -k exponential
python test/test_torch.py -k multinomial
python test/test_torch.py -k kstest

Zeros in 1e8 float16 CPU samples (seed 0, Apple M5):

lambda	before	after
1	3	0
2	3	0
8	24	0

Perf is within noise: float16 exponential_ on 1e7 samples goes 80.4 ms to 81.3 ms (minimum over 15 interleaved runs of the two builds); float64 is unchanged, the compiler folds the check away there.

I have no GPU on hand, so the CUDA change is unverified outside CI.

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @aditew01

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/195077

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit b84d28e with merge base 7fe5288 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[pytorch-bot]

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Skylion007]

Feels like this should be a generic c10::safe_numerics utility maybe?

[qflen]

I added as c10::cast_no_underflow in safe_numerics.h.

It picked up a second caller on the way, since the CUDA kernel's eps/2 floor is applied in accscalar_t and so still narrows to 0 for half once lambda >= 2.