{{ message }}
Reject integral and bool dtypes in uniform decomposition - #196163
Draft
tekinertekin wants to merge 1 commit into
Draft
Reject integral and bool dtypes in uniform decomposition#196163tekinertekin wants to merge 1 commit into
tekinertekin wants to merge 1 commit into
Conversation
torch.rand_like on an integer tensor raises in eager, because uniform_impl_ dispatches over floating types only (AT_DISPATCH_FLOATING_TYPES_AND2 in DistributionTemplates.h). The Python uniform decomposition had no equivalent check, so a traced graph sampled [0, 1) and converted to the input dtype instead: every integral element truncates to 0 and every bool element converts to True. Code asking for random values received a constant, with no error. The check goes in the decomposition, next to the sampling it guards. meta_rand_default repeats it because inductor's replace_random pass rewrites aten.rand before the decomposition runs, so torch.rand(4, dtype=torch.int64) would otherwise still return zeros under inductor. The sibling distributions already reject these dtypes while tracing: randn_like, normal_ and exponential_ all raise, and randint_like is integral by definition, so uniform was the only gap. Fixes pytorch#195673 Test Plan: The new test fails on all four dtypes without the change ("NotImplementedError not raised") and passes with it. ``` python test/test_decomp.py -k test_uniform_integral_dtype python test/test_decomp.py DecompOneOffTestsCPU ``` This PR was authored with AI assistance.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/196163
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #195673
Summary
Calling
torch.rand_likeorTensor.uniform_with integral/bool dtypes raisesNotImplementedErrorin eager mode (viaAT_DISPATCH_FLOATING_TYPES_AND2inDistributionTemplates.h). However, under torch.compile / Inductor, the Pythonuniform()decomposition lacked this check.Traced graphs were silently sampling
[0, 1)floats and casting them to the target dtype:int64,int32,int16,int8,uint8) truncated all values to0.booldtypes converted non-zero floats toTrue.torch.rand(4, dtype=torch.int64)under Inductor also returned zeros becausereplace_randombypasses the decomposition pass.Upstream already notes this gap in
torch/_refs/__init__.py:6639(# TODO: fix inductor rand_like for integer, bool dtypes).This PR adds non-floating/complex dtype checks to both
uniform()decomposition andmeta_rand_default(). Note that this is a behavioral change: code that previously returned zeros silently under Inductor will now raiseNotImplementedError, matching eager mode. Other distributions (randn_like,normal_,exponential_) already raise in both modes, whilerandint_likecontinues to handle integers.Behavior
rand_like(..., dtype=torch.int64)NotImplementedError[0, 0, 0, 0]NotImplementedErrorrand_like(..., dtype=torch.bool)NotImplementedError[True, True, True, True]NotImplementedErrorrand(4, dtype=torch.int64)NotImplementedError[0, 0, 0, 0]NotImplementedErrorAcross 15 tested op/dtype configurations (including float/complex types and invalid integer paths), eager and compiled outputs match 15/15.
Test Plan
Without patch: test_uniform_integral_dtype fails 4/4 (NotImplementedError not raised).
With patch: test_uniform_integral_dtype passes 4/4.
Regression check: DecompOneOffTestsCPU: ran 22, OK (6 skipped, 1 expected failure).
Note: Tests were run against an installed PyTorch 2.14.0 wheel with these exact Python-side patches applied (test_meta.py was skipped locally due to internal API differences with main and relies on CI).
This PR description was generated with AI assistance.