Repository: pytorch
Issue: TorchScript with complex abs doesn't work in backward
Labels: oncall: jit

### 🐛 Describe the bug

Reproduce:
```python
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.bias = torch.nn.Parameter(torch.zeros(1, dtype=torch.float))
        self.act = torch.nn.LeakyReLU()
    
    def forward(self, z):
        out = self.act(torch.abs(z) + self.bias) * torch.exp(1.j * torch.angle(z))
        return out

x = torch.ones(10, dtype=torch.complex64, requires_grad=True)
model = Model()
model = torch.jit.script(model)
# model = torch.compile(model)

for i in range(10):
    out = model(x).sum()
    out.backward()

    print('i = ', i)


```

Stacktrace:
```python
i =  0
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch-test/complex-abs/repro.py", line 21, in <module>
    out.backward()
  File "/home/xwang/Developer/pytorch/torch/_tensor.py", line 484, in backward
    torch.autograd.backward(
  File "/home/xwang/Developer/pytorch/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 183, in <backward op>
        def abs(self):
            def backward(grad_output):
                return grad_output * self.sign()
                                     ~~~~~~~~~ <--- HERE

            return torch.abs(self), backward
RuntimeError: Unlike NumPy, torch.sign is not intended to support complex numbers. Please use torch.sgn instead.
Unlike NumPy, torch.sign is not intended to support complex numbers. Please use torch.sgn instead.
Exception raised from meta at /home/xwang/Developer/pytorch/aten/src/ATen/native/UnaryOps.cpp:276 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xaa (0x7fb1caf2b39a in /home/xwang/Developer/pytorch/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x79 (0x7fb1caef4c3d in /home/xwang/Developer/pytorch/torch/lib/libc10.so)
frame #2: <unknown function> + 0x1700b87 (0x7fb191291b87 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #3: <unknown function> + 0x21caab5 (0x7fb191d5bab5 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #4: <unknown function> + 0x21cab41 (0x7fb191d5bb41 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #5: at::_ops::sign::redispatch(c10::DispatchKeySet, at::Tensor const&) + 0x84 (0x7fb191ab7f64 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #6: <unknown function> + 0x36c075e (0x7fb19325175e in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #7: <unknown function> + 0x36c0c94 (0x7fb193251c94 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #8: <unknown function> + 0x463e417 (0x7fb1941cf417 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #9: <unknown function> + 0x420e8c7 (0x7fb193d9f8c7 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #10: <unknown function> + 0x420fc7d (0x7fb193da0c7d in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #11: <unknown function> + 0x41f1344 (0x7fb193d82344 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #12: <unknown function> + 0x41f2094 (0x7fb193d83094 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #13: <unknown function> + 0x3e55e1b (0x7fb1939e6e1b in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #14: torch::autograd::Engine::evaluate_function(std::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptr<torch::autograd::ReadyQueue> const&) + 0xaeb (0x7fb1939e05db in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #15: torch::autograd::Engine::thread_main(std::shared_ptr<torch::autograd::GraphTask> const&) + 0x63e (0x7fb1939e17be in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #16: torch::autograd::Engine::execute_with_graph_task(std::shared_ptr<torch::autograd::GraphTask> const&, std::shared_ptr<torch::autograd::Node>, torch::autograd::InputBuffer&&) + 0x434 (0x7fb1939dcbf4 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #17: torch::autograd::python::PythonEngine::execute_with_graph_task(std::shared_ptr<torch::autograd::GraphTask> const&, std::shared_ptr<torch::autograd::Node>, torch::autograd::InputBuffer&&) + 0x4b (0x7fb19a2f721b in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
frame #18: torch::autograd::Engine::execute(std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, bool, std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&) + 0x99a (0x7fb1939df42a in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #19: torch::autograd::python::PythonEngine::execute(std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, bool, std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&) + 0x6a (0x7fb19a2f717a in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
frame #20: THPEngine_run_backward(_object*, _object*, _object*) + 0x55d (0x7fb19a2f5c6d in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
<omitting python frames>
frame #37: <unknown function> + 0x23290 (0x7fb1ea3d1290 in /usr/lib/libc.so.6)
frame #38: __libc_start_main + 0x8a (0x7fb1ea3d134a in /usr/lib/libc.so.6)
```

However, if I modify this `self.sign()` to `self.sgn()`, here 
https://github.com/pytorch/pytorch/blob/86269852de0841d5a876fef1570ef0a88797b3a9/torch/csrc/jit/runtime/symbolic_script.cpp#L996-L1001

I got
```python
i =  0
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch-test/complex-abs/repro.py", line 21, in <module>
    out.backward()
  File "/home/xwang/Developer/pytorch/torch/_tensor.py", line 484, in backward
    torch.autograd.backward(
  File "/home/xwang/Developer/pytorch/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Expected isFloatingType(grad.scalar_type()) || (input_is_complex == grad_is_complex) to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.)
Exception raised from validate_outputs at /home/xwang/Developer/pytorch/torch/csrc/autograd/engine.cpp:820 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xaa (0x7f4bd4f5839a in /home/xwang/Developer/pytorch/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x79 (0x7f4bd4f21c3d in /home/xwang/Developer/pytorch/torch/lib/libc10.so)
frame #2: <unknown function> + 0x3e48941 (0x7f4ba37d9941 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #3: torch::autograd::Engine::evaluate_function(std::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptr<torch::autograd::ReadyQueue> const&) + 0x6d0 (0x7f4ba37e01c0 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #4: torch::autograd::Engine::thread_main(std::shared_ptr<torch::autograd::GraphTask> const&) + 0x63e (0x7f4ba37e17be in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #5: torch::autograd::Engine::execute_with_graph_task(std::shared_ptr<torch::autograd::GraphTask> const&, std::shared_ptr<torch::autograd::Node>, torch::autograd::InputBuffer&&) + 0x434 (0x7f4ba37dcbf4 in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #6: torch::autograd::python::PythonEngine::execute_with_graph_task(std::shared_ptr<torch::autograd::GraphTask> const&, std::shared_ptr<torch::autograd::Node>, torch::autograd::InputBuffer&&) + 0x4b (0x7f4baa0f721b in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
frame #7: torch::autograd::Engine::execute(std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, bool, std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&) + 0x99a (0x7f4ba37df42a in /home/xwang/Developer/pytorch/torch/lib/libtorch_cpu.so)
frame #8: torch::autograd::python::PythonEngine::execute(std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, bool, std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&) + 0x6a (0x7f4baa0f717a in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
frame #9: THPEngine_run_backward(_object*, _object*, _object*) + 0x55d (0x7f4baa0f5c6d in /home/xwang/Developer/pytorch/torch/lib/libtorch_python.so)
<omitting python frames>
frame #26: <unknown function> + 0x23290 (0x7f4bfa0ec290 in /usr/lib/libc.so.6)
frame #27: __libc_start_main + 0x8a (0x7f4bfa0ec34a in /usr/lib/libc.so.6)
```

It doesn't work with `torch.compile` either:
```python
Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch/torch/_dynamo/output_graph.py", line 659, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.fake_example_inputs())
  File "/home/xwang/Developer/pytorch/torch/_dynamo/debug_utils.py", line 913, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/__init__.py", line 1153, in __call__
    return self.compile_fn(model_, inputs_)
  File "/home/xwang/Developer/pytorch/torch/_inductor/compile_fx.py", line 398, in compile_fx
    return aot_autograd(
  File "/home/xwang/Developer/pytorch/torch/_dynamo/optimizations/training.py", line 78, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 2333, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
  File "/home/xwang/Developer/pytorch/torch/_dynamo/utils.py", line 90, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 2030, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_tensor_args, aot_config)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 1297, in aot_wrapper_dedupe
    return compiler_fn(flat_fn, leaf_flat_args, aot_config)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 1498, in aot_dispatch_autograd
    fx_g = make_fx(joint_forward_backward, aot_config.decompositions)(
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 704, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/home/xwang/Developer/pytorch/torch/_dynamo/eval_frame.py", line 211, in _fn
    return fn(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 448, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/eval_frame.py", line 211, in _fn
    return fn(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/fx/_symbolic_trace.py", line 739, in trace
    (self.create_arg(fn(*args)),),
  File "/home/xwang/Developer/pytorch/torch/fx/_symbolic_trace.py", line 614, in flatten_fn
    tree_out = root_fn(*tree_args)
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 464, in wrapped
    out = f(*tensors)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 824, in functionalized_joint
    outs = joint_forward_backward(f_primals, f_tangents)
  File "/home/xwang/Developer/pytorch/torch/_functorch/aot_autograd.py", line 791, in joint_forward_backward
    backward_out = torch.autograd.grad(
  File "/home/xwang/Developer/pytorch/torch/autograd/__init__.py", line 266, in grad
    return handle_torch_function(
  File "/home/xwang/Developer/pytorch/torch/overrides.py", line 1520, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/home/xwang/Developer/pytorch/torch/_inductor/overrides.py", line 37, in __torch_function__
    return func(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/autograd/__init__.py", line 300, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 491, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 516, in inner_torch_dispatch
    out = proxy_call(self, func, args, kwargs)
  File "/home/xwang/Developer/pytorch/torch/fx/experimental/proxy_tensor.py", line 352, in proxy_call
    out = func(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_ops.py", line 284, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: aten::_conj() Expected a value of type 'Tensor' for argument 'self' but instead found type 'complex'.
Position: 0
Value: 1j
Declaration: aten::_conj(Tensor(a) self) -> Tensor(a)
Cast error details: Unable to cast 1j to Tensor

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/xwang/Developer/pytorch-test/complex-abs/repro.py", line 20, in <module>
    out = model(x).sum()
  File "/home/xwang/Developer/pytorch/torch/nn/modules/module.py", line 1482, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/eval_frame.py", line 82, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/eval_frame.py", line 211, in _fn
    return fn(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/eval_frame.py", line 332, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/convert_frame.py", line 480, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/convert_frame.py", line 103, in _fn
    return fn(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/utils.py", line 90, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/convert_frame.py", line 339, in _convert_frame_assert
    return _compile(
  File "/home/xwang/Developer/pytorch/torch/_dynamo/convert_frame.py", line 400, in _compile
    out_code = transform_code_object(code, transform)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/bytecode_transformation.py", line 341, in transform_code_object
    transformations(instructions, code_options)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/convert_frame.py", line 387, in transform
    tracer.run()
  File "/home/xwang/Developer/pytorch/torch/_dynamo/symbolic_convert.py", line 1687, in run
    super().run()
  File "/home/xwang/Developer/pytorch/torch/_dynamo/symbolic_convert.py", line 538, in run
    and self.step()
  File "/home/xwang/Developer/pytorch/torch/_dynamo/symbolic_convert.py", line 501, in step
    getattr(self, inst.opname)(inst)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/symbolic_convert.py", line 1753, in RETURN_VALUE
    self.output.compile_subgraph(self)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/output_graph.py", line 511, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/output_graph.py", line 583, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/xwang/Developer/pytorch/torch/_dynamo/output_graph.py", line 664, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e) from e
torch._dynamo.exc.BackendCompilerFailed: debug_wrapper raised RuntimeError: aten::_conj() Expected a value of type 'Tensor' for argument 'self' but instead found type 'complex'.
Position: 0
Value: 1j
Declaration: aten::_conj(Tensor(a) self) -> Tensor(a)
Cast error details: Unable to cast 1j to Tensor

Set torch._dynamo.config.verbose=True for more information


You can suppress this exception and fall back to eager by setting:
    torch._dynamo.config.suppress_errors = True

```

### Versions

1e347b737bbb7666dc240bb94b4c60403ad882ae
cuda 11.8

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel @ptrblck @ngimel