Repository: pytorch/torchdynamo

Issue:
torch.cat: Does not support different dtypes for operands
This is supported by Eager PyTorch, see minimal repro:


```
import torchvision
import torch
import torchdynamo

dtype = torch.float16

a = torch.randn((1,), device="cuda", dtype=torch.float32)
b = torch.randn((1,), device="cuda", dtype=torch.float16)

def simple_op():
    return torch.cat([a, b], dim=0)

@torchdynamo.optimize("inductor")
def run_once_dynamo_ind():
    return simple_op()

y = simple_op() # Working fine
print("Eager Pytorch returned tensor", y)
run_once_dynamo_ind()
```

Output:
```
WARNING:torchinductor.lowering:make_fallback(aten.cumsum): a decomposition exists, we should switch to it
WARNING:torchinductor.lowering:make_fallback(aten.unfold): a decomposition exists, we should switch to it
WARNING:torchinductor.lowering:make_fallback(aten.unfold_backward): a decomposition exists, we should switch to it
Eager Pytorch returned tensor tensor([-1.6989, -1.2197], device='cuda:0')
Traceback (most recent call last):
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/graph.py", line 257, in call_function
    return lowerings[target](*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/lowering.py", line 193, in wrapped
    return decomp_fn(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/lowering.py", line 746, in cat
    return TensorBox(ir.ConcatKernel.create(inputs, dim))
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/ir.py", line 2135, in create
    assert inputs[i].get_dtype() == dtype
AssertionError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_ti.py", line 19, in <module>
    run_once_dynamo_ind()
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/eval_frame.py", line 175, in _fn
    return fn(*args, **kwargs)
  File "test_ti.py", line 13, in run_once_dynamo_ind
    @torchdynamo.optimize("inductor")
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/eval_frame.py", line 175, in _fn
    return fn(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 797, in forward
    return compiled_f(
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 783, in new_func
    compiled_fn = create_aot_dispatcher_function(
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 512, in create_aot_dispatcher_function
    return aot_dispatch_base(flat_fn, fake_flat_tensor_args, aot_config)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 282, in aot_dispatch_base
    compiled_fw = aot_config.fw_compiler(fw_module, flat_args)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/functorch/_src/aot_autograd.py", line 238, in f
    out_f = compiler(fx_g, inps)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/compile_fx.py", line 333, in fw_compiler
    return compile_fx_inner(
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/debug_utils.py", line 445, in debug_wrapper
    compiled_fn = compiler_fn(gm, example_inputs, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/debug.py", line 180, in inner
    return fn(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/compile_fx.py", line 102, in compile_fx_inner
    graph.run(*example_inputs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchdynamo/utils.py", line 76, in time_wrapper
    r = func(*args, **kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/graph.py", line 146, in run
    return super().run(*args)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torch/fx/interpreter.py", line 130, in run
    self.env[node] = self.run_node(node)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/graph.py", line 314, in run_node
    result = super().run_node(n)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torch/fx/interpreter.py", line 171, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/scratch/XXXX/xformers-py38-pytorch113-cu116/lib/python3.8/site-packages/torchinductor/graph.py", line 259, in call_function
    raise LoweringException(e, target, args, kwargs) from e
torchinductor.exc.LoweringException: AssertionError: 
  target: aten.cat.default
  args[0]: [TensorBox(StorageBox(
    InputBuffer(name='arg0_1', layout=FixedLayout('cuda', torch.float32, size=[1], stride=[1]))
  )), TensorBox(StorageBox(
    InputBuffer(name='arg1_1', layout=FixedLayout('cuda', torch.float16, size=[1], stride=[1]))
  ))]

While executing %cat : [#users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%arg0_1, %arg1_1],), kwargs = {})
Original traceback:
Module stack: {}
  File "test_ti.py", line 11, in simple_op
    return torch.cat([a, b], dim=0)
 |   File "test_ti.py", line 15, in run_once_dynamo_ind
    return simple_op()

```


Discussion:

Resolving patch:
diff --git a/torchinductor/lowering.py b/torchinductor/lowering.py
--- a/torchinductor/lowering.py
+++ b/torchinductor/lowering.py
@@ -157,6 +157,11 @@ def _register_lowering(
     @functools.wraps(decomp_fn)
     def wrapped(*args, **kwargs):
         args = list(args)
+        unpacked = False
+        # TODO maybe we need to use pytrees here
+        if len(args) == 1 and isinstance(args[0], (list, tuple)):
+            unpacked = True
+            args = args[0]
         # Only look at args that are Tensors
         indices = [i for i, x in enumerate(args) if isinstance(x, TensorBox)]
         # kwargs tensors not supported yet
@@ -173,14 +178,20 @@ def wrapped(*args, **kwargs):
                 dtype = get_promoted_dtype(
                     *promoting_args, type_promotion_kind=type_promotion_kind
                 )
-            for i in indices:
-                args[i] = to_dtype(args[i], dtype)
+            # sometimes args are an immutable list so we can't mutate them
+            new_args = []
             for i in range(len(args)):
-                if isinstance(args[i], ir.Constant):
-                    args[i] = ir.Constant(
-                        args[i].value, dtype, args[indices[0]].get_device()
+                if i in indices:
+                    new_args.append(to_dtype(args[i], dtype))
+                elif isinstance(args[i], ir.Constant):
+                    new_args.append(
+                        ir.Constant(args[i].value, dtype, args[indices[0]].get_device())
                     )
-
+                else:
+                    new_args.append(args[i])
+            args = new_args
+        if unpacked:
+            args = [args]
         if broadcast and indices:
             for i, x in zip(indices, broadcast_tensors(*[args[i] for i in indices])):
                 args[i] = x


Test patch:
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
--- a/test/inductor/test_torchinductor.py
+++ b/test/inductor/test_torchinductor.py
@@ -1995,7 +1995,11 @@ def fn(x):
     def test_cat(self):
         def fn(a):
             tmp = a * 2
-            return torch.cat((a, a[:, :4] + 1, a + 2), -1), torch.cat((tmp, tmp), 0)
+            return (
+                torch.cat((a, a[:, :4] + 1, a + 2), -1),
+                torch.cat((tmp, tmp), 0),
+                torch.cat((tmp, tmp.double()), 0),
+            )
 
         self.common(
             fn,