Coverage for nondeterministic_seeded, respect it in constant prop (#83650)
- nondeterministic_seeded was not applied to enough functions. I added
some heuristics to codegen for identifying functions that are likely
to be random and added a bunch of these tags to functions. Not sure
I got all of them.
- Don't constant propagate through nondeterministic functions in FX
tracing.
It would be better to do some testing for the tag but this would be quite an effort.
Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83650
Approved by: https://github.com/bdhirsh, https://github.com/eellison
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index af744e1..6db18f2 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -256,18 +256,25 @@
tags: nondeterministic_seeded
- func: dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: feature_dropout(Tensor input, float p, bool train) -> Tensor
+ tags: nondeterministic_seeded
- func: feature_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: alpha_dropout(Tensor input, float p, bool train) -> Tensor
+ tags: nondeterministic_seeded
- func: alpha_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: feature_alpha_dropout(Tensor input, float p, bool train) -> Tensor
+ tags: nondeterministic_seeded
- func: feature_alpha_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: abs(Tensor self) -> Tensor
device_check: NoCheck # TensorIterator
@@ -928,6 +935,7 @@
- func: bernoulli.out(Tensor self, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: function
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: bernoulli_out
MPS: bernoulli_out_mps
@@ -935,6 +943,7 @@
- func: bernoulli_.Tensor(Tensor(a!) self, Tensor p, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: method
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: bernoulli_
MPS: bernoulli_mps_
@@ -943,6 +952,7 @@
- func: bernoulli_.float(Tensor(a!) self, float p=0.5, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: method
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: bernoulli_
MPS: bernoulli_mps_
@@ -4007,10 +4017,12 @@
dispatch:
CompositeExplicitAutograd: rand
autogen: rand.names_out
+ tags: nondeterministic_seeded
- func: rand.generator_with_names(int[] size, *, Generator? generator, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
device_check: NoCheck
device_guard: False
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: rand
autogen: rand.generator_with_names_out
@@ -4021,14 +4033,17 @@
CompositeExplicitAutograd: rand
- func: rand.generator(int[] size, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: rand
- func: rand.out(int[] size, *, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: rand_out
- func: rand.generator_out(int[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: rand_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
tags: nondeterministic_seeded
@@ -4044,6 +4059,7 @@
CompositeExplicitAutograd: randint
- func: randint.generator(int high, int[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint
@@ -4053,22 +4069,27 @@
CompositeExplicitAutograd: randint
- func: randint.low_generator(int low, int high, int[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint
- func: randint.out(int high, int[] size, *, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint_out
- func: randint.generator_out(int high, int[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint_out
- func: randint.low_out(int low, int high, int[] size, *, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint_out
- func: randint.low_generator_out(int low, int high, int[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randint_out
@@ -4094,10 +4115,12 @@
CompositeExplicitAutograd: randn
- func: randn.generator(int[] size, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randn
- func: randn.names(int[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
device_check: NoCheck
device_guard: False
dispatch:
@@ -4105,6 +4128,7 @@
autogen: randn.names_out
- func: randn.generator_with_names(int[] size, *, Generator? generator, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
device_check: NoCheck
device_guard: False
dispatch:
@@ -4112,8 +4136,10 @@
autogen: randn.generator_with_names_out
- func: randn.out(int[] size, *, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: randn.generator_out(int[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
- func: randn_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
tags: nondeterministic_seeded
@@ -4129,14 +4155,17 @@
CompositeExplicitAutograd: randperm
- func: randperm.generator(int n, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randperm
- func: randperm.out(int n, *, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: randperm_out
- func: randperm.generator_out(int n, *, Generator? generator, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CPU: randperm_out_cpu
CUDA: randperm_out_cuda
@@ -4324,6 +4353,7 @@
tags: nondeterministic_seeded
- func: rrelu_(Tensor(a!) self, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!)
+ tags: nondeterministic_seeded
device_check: NoCheck # TensorIterator
- func: relu(Tensor self) -> Tensor
@@ -5622,6 +5652,7 @@
autogen: _dirichlet_grad.out
- func: _sample_dirichlet(Tensor self, Generator? generator=None) -> Tensor
+ tags: nondeterministic_seeded
variants: function
dispatch:
CPU: _s_dirichlet_cpu
@@ -7492,6 +7523,7 @@
- func: random_.from(Tensor(a!) self, int from, int? to, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: method
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: random_
Meta: random_meta_
@@ -7500,6 +7532,7 @@
- func: random_.to(Tensor(a!) self, int to, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: random_
@@ -7509,6 +7542,7 @@
- func: random_(Tensor(a!) self, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: random_
@@ -7517,6 +7551,7 @@
- func: uniform_(Tensor(a!) self, float from=0, float to=1, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: uniform_
@@ -7527,12 +7562,14 @@
- func: cauchy_(Tensor(a!) self, float median=0, float sigma=1, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
variants: method
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: cauchy_
autogen: cauchy, cauchy.out
- func: log_normal_(Tensor(a!) self, float mean=1, float std=2, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: log_normal_
@@ -7540,6 +7577,7 @@
- func: exponential_(Tensor(a!) self, float lambd=1, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: exponential_
@@ -7548,6 +7586,7 @@
- func: geometric_(Tensor(a!) self, float p, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: geometric_
@@ -8267,6 +8306,7 @@
# TODO: remove dispatch section when porting TH CUDA to ATen
- func: multinomial.out(Tensor self, int num_samples, bool replacement=False, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: multinomial_out
@@ -8960,6 +9000,7 @@
- func: normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!)
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
variants: method
dispatch:
CPU, CUDA: normal_
@@ -8973,10 +9014,12 @@
# but we can't due to overload ambiguity with normal.Tensor_float.
- func: normal_functional(Tensor self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor
device_check: NoCheck # TensorIterator
+ tags: nondeterministic_seeded
dispatch:
CompositeExplicitAutograd: normal_functional
- func: normal.Tensor_float_out(Tensor mean, float std=1, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!)
+ tags: nondeterministic_seeded
dispatch:
CPU, CUDA: normal_out
MPS: normal_mps_out
@@ -8994,6 +9037,7 @@
CPU, CUDA: normal_out
Meta: normal_out_meta
MPS: normal_mps_out
+ tags: nondeterministic_seeded
- func: normal.float_Tensor(float mean, Tensor std, *, Generator? generator=None) -> Tensor
dispatch:
@@ -9007,6 +9051,7 @@
CPU, CUDA: normal_out
Meta: normal_out_meta
MPS: normal_mps_out
+ tags: nondeterministic_seeded
- func: normal.Tensor_Tensor(Tensor mean, Tensor std, *, Generator? generator=None) -> Tensor
dispatch:
@@ -9018,10 +9063,12 @@
- func: normal.float_float(float mean, float std, int[] size, *, Generator? generator=None, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
dispatch:
CompositeExplicitAutograd: normal
+ tags: nondeterministic_seeded
- func: normal.float_float_out(float mean, float std, int[] size, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!)
dispatch:
CompositeExplicitAutograd: normal_out
+ tags: nondeterministic_seeded
- func: alias(Tensor(a) self) -> Tensor(a)
variants: method, function
@@ -10272,6 +10319,7 @@
- func: rrelu_with_noise.out(Tensor self, Tensor noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None, *, Tensor(a!) out) -> Tensor(a!)
python_module: nn
+ tags: nondeterministic_seeded
dispatch:
CPU: rrelu_with_noise_out_cpu
CUDA: rrelu_with_noise_out_cuda
@@ -10291,6 +10339,7 @@
- func: rrelu_with_noise_(Tensor(a!) self, Tensor noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!)
python_module: nn
+ tags: nondeterministic_seeded
dispatch:
CPU: rrelu_with_noise_cpu_
CUDA: rrelu_with_noise_cuda_
diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
index c3bda7e..0dc46da 100644
--- a/test/test_proxy_tensor.py
+++ b/test/test_proxy_tensor.py
@@ -454,6 +454,17 @@
lambda: make_fx(f, tracing_mode=self.tracing_mode)()
)
+ def test_constant_random(self):
+ def f():
+ val = torch.tensor([2.0])
+ val.normal_()
+ return val.item()
+
+ self.assertRaisesRegex(
+ RuntimeError, "data-dependent",
+ lambda: make_fx(f, tracing_mode=self.tracing_mode)()
+ )
+
def test_decomposition_interpreter(self):
def fn(x):
return torch.nn.functional.silu(x)
diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py
index 991037d..a8b9c7e 100644
--- a/torch/fx/experimental/proxy_tensor.py
+++ b/torch/fx/experimental/proxy_tensor.py
@@ -240,7 +240,8 @@
constant = None
# NB: do NOT include factories as constants
if (
- all_constant
+ torch.Tag.nondeterministic_seeded not in func_overload.tags # type: ignore[attr-defined]
+ and all_constant
and any_constant
and pytree.tree_all_only(torch.Tensor, lambda t: t.numel() <= CONSTANT_NUMEL_LIMIT, out)
):
diff --git a/torchgen/model.py b/torchgen/model.py
index 76bf473..ecf34db 100644
--- a/torchgen/model.py
+++ b/torchgen/model.py
@@ -882,6 +882,20 @@
"device_check not allowed to be enabled"
)
+ # NB: if your function accidentally has rand/dropout/... in its name
+ # but is not actually random, feel free to amend this to special case
+ if (
+ "rand" in str(self.func.name)
+ or (
+ "dropout" in str(self.func.name)
+ # Backwards of dropout is typically deterministic
+ and "backward" not in str(self.func.name)
+ and str(self.func.name.name) not in ["_cudnn_init_dropout_state"]
+ )
+ or self.func.arguments.has_generator_arg()
+ ):
+ assert "nondeterministic_seeded" in self.tags, str(self.func.name)
+
@property
def has_composite_kernel(self) -> bool:
return (
@@ -1668,6 +1682,9 @@
def is_tensor_like(self) -> bool:
raise NotImplementedError
+ def is_generator_like(self) -> bool:
+ raise NotImplementedError
+
def is_nullable(self) -> bool:
raise NotImplementedError
@@ -1714,6 +1731,9 @@
def is_tensor_like(self) -> bool:
return self.name == BaseTy.Tensor
+ def is_generator_like(self) -> bool:
+ return self.name == BaseTy.Generator
+
def is_nullable(self) -> bool:
return False
@@ -1737,6 +1757,9 @@
def is_tensor_like(self) -> bool:
return self.elem.is_tensor_like()
+ def is_generator_like(self) -> bool:
+ return self.elem.is_generator_like()
+
def is_nullable(self) -> bool:
return True
@@ -1764,6 +1787,9 @@
"""
return False
+ def is_generator_like(self) -> bool:
+ return False
+
def is_nullable(self) -> bool:
"""
Assume a custom class is not nullable.
@@ -1796,6 +1822,9 @@
def is_tensor_like(self) -> bool:
return self.elem.is_tensor_like()
+ def is_generator_like(self) -> bool:
+ return self.elem.is_generator_like()
+
def is_nullable(self) -> bool:
return self.elem.is_nullable()
@@ -2101,6 +2130,9 @@
def has_tensor_arg(self) -> bool:
return any(a.type.is_tensor_like() for a in self.flat_non_out)
+ def has_generator_arg(self) -> bool:
+ return any(a.type.is_generator_like() for a in self.flat_non_out)
+
def signature(self, *, strip_default: bool = False) -> "Arguments":
# dataclasses.replace could be used here, but it is less
# type safe so for now I've opted to type everything out
diff --git a/torchgen/native_function_generation.py b/torchgen/native_function_generation.py
index 6ff25b0..70755d4 100644
--- a/torchgen/native_function_generation.py
+++ b/torchgen/native_function_generation.py
@@ -340,7 +340,7 @@
has_composite_explicit_autograd_non_functional_kernel=False,
# Every generated NativeFunction gets a "generated" tag, so it's easy to tell
# which NativeFunction objects did not come directly from native_functions.yaml.
- tags=set(["generated"]),
+ tags=set(["generated"]) | (f.tags & {"nondeterministic_seeded"}),
namespace=f.namespace,
),
backend_metadata,
