`torch.compile` fails with `jacfwd` when multiplying/dividing float and tensor #125078

cw-tan · 2024-04-26T22:12:45Z

🐛 Describe the bug

The following minimal example fails

import torch
from torch.func import jacfwd

def func(x):
    two = 2.0
    return two * x


def jac_func(x):
    return jacfwd(func, argnums=(0,))(x)


compiled_jac_func = torch.compile(jac_func)
compiled_jac_func(torch.ones((3,), dtype=torch.float64))

with following last few lines in the error message (longer error log below).

torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in function mul>(*(2.0, GradTrackingTensor(lvl=2, value=
    FakeTensor(..., size=(3,), dtype=torch.float64)
)), **{}):
aten::alias() Expected a value of type 'Tensor' for argument 'self' but instead found type 'float'.
Position: 0
Value: 2.0
Declaration: aten::alias(Tensor(a) self) -> Tensor(a)
Cast error details: Unable to cast 2.0 to Tensor

from user code:
   File "/home/cwtan/florch/tests/test.py", line 10, in jac_func
    return jacfwd(func, argnums=(0,))(x)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1312, in wrapper_fn
    results = vmap(push_jvp, randomness=randomness)(basis)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/apis.py", line 200, in wrapped
    return vmap_impl(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
    return _flat_vmap(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 48, in fn
    return f(*args, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 480, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1301, in push_jvp
    output = _jvp_with_argnums(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 48, in fn
    return f(*args, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1141, in _jvp_with_argnums
    result_duals = func(*duals)
  File "/home/cwtan/florch/tests/test.py", line 6, in func
    return two * x

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

Similar error if it's return x / two. Trying instead

import torch
from torch.func import jacfwd

def func(x):
    two = torch.tensor([2.0], dtype=x.dtype, device=x.device)
    return two * x


def jac_func(x):
    return jacfwd(func, argnums=(0,))(x)


compiled_jac_func = torch.compile(jac_func)
compiled_jac_func(torch.ones((3,), dtype=torch.float64))

also fails with a different NotImplementedError (full error below).

torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
NotImplementedError: Cannot access storage of TensorWrapper

Error logs

For the first example where two = 2.0.

refreshing <module 'torch.ops.quantized_decomposed' from 'torch.ops'> quantize_per_tensor
refreshing <module 'torch.ops.quantized_decomposed' from 'torch.ops'> quantize_per_tensor
refreshing <module 'torch.ops.quantized_decomposed' from 'torch.ops'> dequantize_per_tensor
refreshing <module 'torch.ops.quantized_decomposed' from 'torch.ops'> dequantize_per_tensor
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] failed while attempting to run meta for aten.alias.default
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] Traceback (most recent call last):
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]   File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1493, in _dispatch_impl
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]     r = func(*args, **kwargs)
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]         ^^^^^^^^^^^^^^^^^^^^^
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]   File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_ops.py", line 630, in __call__
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]     return self_._op(*args, **kwargs)
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] RuntimeError: aten::alias() Expected a value of type 'Tensor' for argument 'self' but instead found type 'float'.
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] Position: 0
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] Value: 2.0
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] Declaration: aten::alias(Tensor(a) self) -> Tensor(a)
E0426 18:03:00.287000 140386476729408 torch/_subclasses/fake_tensor.py:1497] [0/0] Cast error details: Unable to cast 2.0 to Tensor
Traceback (most recent call last):
  File "/home/cwtan/florch/tests/test.py", line 14, in <module>
    compiled_jac_func(torch.ones((3,), dtype=torch.float64))
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 403, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 977, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 818, in _convert_frame
    result = inner_convert(
             ^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 411, in _convert_frame_assert
    return _compile(
           ^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_utils_internal.py", line 70, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 700, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 268, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 568, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1116, in transform_code_object
    transformations(instructions, code_options)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 173, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 515, in transform
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2237, in run
    super().run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1848, in CALL
    self.call_function(fn, args, kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1848, in CALL
    self.call_function(fn, args, kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/higher_order_ops.py", line 1202, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1848, in CALL
    self.call_function(fn, args, kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 492, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1301, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 730, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 736, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2418, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2534, in inline_call_
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1812, in BINARY_OP
    return _binary_op_lookup[inst.arg](self, inst)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 229, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/builtin.py", line 946, in call_function
    return handler(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/builtin.py", line 925, in _handle_insert_op_in_graph
    return wrap_fx_proxy(tx, proxy)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1434, in wrap_fx_proxy
    return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1519, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1821, in get_fake_value
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1753, in get_fake_value
    ret_val = wrap_fake_exception(
              ^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1268, in wrap_fake_exception
    return fn()
           ^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1754, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1889, in run_node
    raise RuntimeError(make_error_message(e)).with_traceback(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1871, in run_node
    return node.target(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 904, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1262, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 987, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_subclasses/fake_tensor.py", line 1493, in _dispatch_impl
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_ops.py", line 630, in __call__
    return self_._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.TorchRuntimeError: Failed running call_function <built-in function mul>(*(2.0, GradTrackingTensor(lvl=2, value=
    FakeTensor(..., size=(3,), dtype=torch.float64)
)), **{}):
aten::alias() Expected a value of type 'Tensor' for argument 'self' but instead found type 'float'.
Position: 0
Value: 2.0
Declaration: aten::alias(Tensor(a) self) -> Tensor(a)
Cast error details: Unable to cast 2.0 to Tensor

from user code:
   File "/home/cwtan/florch/tests/test.py", line 10, in jac_func
    return jacfwd(func, argnums=(0,))(x)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1312, in wrapper_fn
    results = vmap(push_jvp, randomness=randomness)(basis)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/apis.py", line 200, in wrapped
    return vmap_impl(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 331, in vmap_impl
    return _flat_vmap(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 48, in fn
    return f(*args, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 480, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1301, in push_jvp
    output = _jvp_with_argnums(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 48, in fn
    return f(*args, **kwargs)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1141, in _jvp_with_argnums
    result_duals = func(*duals)
  File "/home/cwtan/florch/tests/test.py", line 6, in func
    return two * x

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

For the second example where two = torch.tensor([2.0], dtype=x.dtype, device=x.device).

Traceback (most recent call last):
  File "/home/cwtan/florch/tests/test.py", line 14, in <module>
    compiled_jac_func(torch.ones((3,), dtype=torch.float64))
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 403, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 977, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 818, in _convert_frame
    result = inner_convert(
             ^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 411, in _convert_frame_assert
    return _compile(
           ^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_utils_internal.py", line 70, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 700, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 268, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 568, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1116, in transform_code_object
    transformations(instructions, code_options)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 173, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 515, in transform
    tracer.run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2237, in run
    super().run()
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2394, in RETURN_VALUE
    self._return(inst)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2379, in _return
    self.output.compile_subgraph(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1091, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1283, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 268, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1374, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1355, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 127, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/__init__.py", line 1744, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1434, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 65, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 958, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 268, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 558, in create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/cwtan/anaconda3/envs/florch/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/collect_metadata_analysis.py", line 262, in inner
    curr_storage = StorageWeakRef(o.untyped_storage())
                                  ^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
NotImplementedError: Cannot access storage of TensorWrapper

Minified repro

No response

Versions

Collecting environment information...
PyTorch version: 2.4.0.dev20240425+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.146.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A1000 6GB Laptop GPU
Nvidia driver version: 537.77
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             20
On-line CPU(s) list:                0-19
Vendor ID:                          GenuineIntel
Model name:                         13th Gen Intel(R) Core(TM) i7-13800H
CPU family:                         6
Model:                              186
Thread(s) per core:                 2
Core(s) per socket:                 10
Socket(s):                          1
Stepping:                           2
BogoMIPS:                           5836.80
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves avx_vnni umip waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          480 KiB (10 instances)
L1i cache:                          320 KiB (10 instances)
L2 cache:                           12.5 MiB (10 instances)
L3 cache:                           24 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] flake8==7.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] pytorch-triton==3.0.0+45fff310c8
[pip3] torch==2.4.0.dev20240425+cu121
[pip3] torch-nl==0.3
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch-triton            3.0.0+45fff310c8          pypi_0    pypi
[conda] torch                     2.4.0.dev20240425+cu121          pypi_0    pypi
[conda] torch-nl                  0.3                      pypi_0    pypi

cc @zou3519 @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

The text was updated successfully, but these errors were encountered:

ayush0x00 · 2024-04-27T14:50:27Z

Well, it seems to be working on torch==2.3.0.dev20240212 on mps.
But torch==2.0.1+cu117 is giving error RuntimeError: Cannot access data pointer of Tensor that doesn't have storage

jbschlosser · 2024-04-29T20:14:13Z

Possibly related: #124731

cw-tan added the oncall: pt2 label Apr 26, 2024

jbschlosser added triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module module: vmap module: dynamo labels Apr 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

`torch.compile` fails with `jacfwd` when multiplying/dividing float and tensor #125078

`torch.compile` fails with `jacfwd` when multiplying/dividing float and tensor #125078

cw-tan commented Apr 26, 2024 •

edited by pytorch-bot bot

ayush0x00 commented Apr 27, 2024

jbschlosser commented Apr 29, 2024

torch.compile fails with jacfwd when multiplying/dividing float and tensor #125078

torch.compile fails with jacfwd when multiplying/dividing float and tensor #125078

Comments

cw-tan commented Apr 26, 2024 • edited by pytorch-bot bot

🐛 Describe the bug

Error logs

Minified repro

Versions

ayush0x00 commented Apr 27, 2024

jbschlosser commented Apr 29, 2024

`torch.compile` fails with `jacfwd` when multiplying/dividing float and tensor #125078

`torch.compile` fails with `jacfwd` when multiplying/dividing float and tensor #125078

cw-tan commented Apr 26, 2024 •

edited by pytorch-bot bot