aten::rrelu_with_noise function mutates noise parameter

[ttanpcs]

🐛 Describe the bug

class TestMode(TorchDispatchMode):
        def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                print(func._schema.name)
                print(args)
                out = func(*args, **kwargs)
                print(args)
                return out

a = torch.rand((3,3))
with enable_torch_dispatch_mode(TestMode()):
        torch.nn.functional.rrelu(a, training=True)

returns
aten::empty_like
(tensor([[0.0290, 0.4019, 0.2598],
[0.3666, 0.0583, 0.7006],
[0.0518, 0.4681, 0.6738]]),)
(tensor([[0.0290, 0.4019, 0.2598],
[0.3666, 0.0583, 0.7006],
[0.0518, 0.4681, 0.6738]]),)
aten::rrelu_with_noise
(tensor([[0.0290, 0.4019, 0.2598],
[0.3666, 0.0583, 0.7006],
[0.0518, 0.4681, 0.6738]]), tensor([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]), 0.125, 0.3333333333333333, True)
(tensor([[0.0290, 0.4019, 0.2598],
[0.3666, 0.0583, 0.7006],
[0.0518, 0.4681, 0.6738]]), tensor([[1., 1., 1.],
[1., 1., 1.],
[1., 1., 1.]]), 0.125, 0.3333333333333333, True)

^ Where the noise parameter mutates before and after aten::rrelu_with_noise. However the function schema in native_functions.yaml doesn't specify that the noise parameter is mutable.

Versions

N/A

[bdhirsh]

Nice catch. One thing to check would probably be: is at::rrelu_with_noise a "running-average" style operator, like batch norm? From the example in the description, it looks like noise isn't passed by the user; so is it stored as some internal state?

If it is, then rrelu_with_noise should be classified as "mutable" op; it's not technically an inplace or out= op (those both exist separately), but at least one of the inputs that it takes in gets mutated.

If the above is true, then i imagine the final set of schemas in native_functions.yaml should look something like:

- func: rrelu_with_noise.out(Tensor self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None, *, Tensor(a!) out) -> Tensor(a!)
  python_module: nn
  dispatch:
    CPU: rrelu_with_noise_out_cpu
    CUDA: rrelu_with_noise_out_cuda

- func: rrelu_with_noise(Tensor self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor
  python_module: nn
  dispatch:
    CPU: rrelu_with_noise_cpu
    CUDA: rrelu_with_noise_cuda
  tags: nondeterministic_seeded
  # There should be a "real", functional version of rrelu_with_noise that returns the updated noise instead of mutating it
  # It should just be auto-generated through codegen (there's some existing infra that can do this)
  autogen: rrelu_with_noise.functional

- func: rrelu_with_noise_(Tensor(a!) self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!)
  python_module: nn
  dispatch:
    CPU: rrelu_with_noise_cpu_
    CUDA: rrelu_with_noise_cuda_