Repository: pytorch
Issue: Compilation Fails with torch.sparse and "fullgraph=True"
Labels: module: sparse,triaged,oncall: pt2

### 🐛 Describe the bug

When you have a function that uses torch.sparse, torch.compile can compile it, but not with "fullgraph=True" 

### Error logs

(e3nn_env) aglover@login36:~/e3nn_benchmarking> python minimal_sparse_compilation_error.py 
/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py:21: UserWarning: Sparse CSR tensor support is in beta state. If you miss a functionality in the sparse tensor support, please submit a feature request to https://github.com/pytorch/pytorch/issues. (Triggered internally at /home/conda/feedstock_root/build_artifacts/libtorch_1718580740865/work/aten/src/ATen/SparseCsrTensorImpl.cpp:53.)
  z = torch.randn(size=(3,9), dtype=torch.float32, device=device).to_sparse_csr()
0.000072
I0713 17:21:46.222000 140632849619840 torch/_dynamo/logging.py:55] [0/0] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34
I0713 17:21:46.233000 140632849619840 torch/_dynamo/convert_frame.py:538] [0/0] Restarting analysis due to _dynamo/symbolic_convert.py:142 in fail_and_restart_analysis
I0713 17:21:46.234000 140632849619840 torch/_dynamo/logging.py:55] [0/0_1] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] WON'T CONVERT inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py line 34
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== TorchDynamo Stack Trace ==========
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last):
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = inner_convert(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return _compile(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return func(*args, **kwds)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     guarded_code = compile_inner(code, one_graph, hooks, transform)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     r = func(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     out_code = transform_code_object(code, transform)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     transformations(instructions, code_options)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     tracer.run()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     super().run()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     and self.step()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     getattr(self, inst.opname)(inst)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 484, in wrapper
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return handle_graph_break(self, inst, speculation.reason)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 548, in handle_graph_break
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.output.compile_subgraph(self, reason=reason)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 961, in compile_subgraph
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     and all(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 962, in <genexpr>
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     not isinstance(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in __instancecheck__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     instance = instance.realize()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self._cache.realize(self.parents_tracker)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.vt = VariableBuilder(tx, self.source)(self.value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     keywords[k] = VariableBuilder(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return type_dispatch(self, value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     unimplemented(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     raise Unsupported(msg)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code:
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]    File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting:
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     import torch._dynamo
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     torch._dynamo.config.suppress_errors = True
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== The above exception occurred while processing the following code ==========
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in <module>
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     main()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in main
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in print_performance
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in <listcomp>
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = model(*example_inputs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in <lambda>
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] ==========
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last):
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = inner_convert(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return _compile(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return func(*args, **kwds)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     guarded_code = compile_inner(code, one_graph, hooks, transform)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     r = func(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     out_code = transform_code_object(code, transform)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     transformations(instructions, code_options)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     tracer.run()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     super().run()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     and self.step()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     getattr(self, inst.opname)(inst)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 484, in wrapper
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return handle_graph_break(self, inst, speculation.reason)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 548, in handle_graph_break
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.output.compile_subgraph(self, reason=reason)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 961, in compile_subgraph
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     and all(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 962, in <genexpr>
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     not isinstance(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in __instancecheck__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     instance = instance.realize()
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self._cache.realize(self.parents_tracker)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.vt = VariableBuilder(tx, self.source)(self.value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     keywords[k] = VariableBuilder(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return type_dispatch(self, value)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     unimplemented(
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     raise Unsupported(msg)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code:
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]    File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting:
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     import torch._dynamo
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822]     torch._dynamo.config.suppress_errors = True
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.243000 140632849619840 torch/_dynamo/logging.py:55] [1/0] Step 1: torchdynamo start tracing func_with_sparse /global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py:6
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] WON'T CONVERT func_with_sparse /global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py line 6
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== TorchDynamo Stack Trace ==========
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last):
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = inner_convert(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return _compile(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return func(*args, **kwds)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     guarded_code = compile_inner(code, one_graph, hooks, transform)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     r = func(*args, **kwargs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     out_code = transform_code_object(code, transform)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     transformations(instructions, code_options)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     tracer.run()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     super().run()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     and self.step()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     getattr(self, inst.opname)(inst)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 242, in impl
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.push(fn_var.call_function(self, self.popn(nargs), {}))
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in call_function
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     args = [v.realize() for v in args]
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in <listcomp>
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     args = [v.realize() for v in args]
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self._cache.realize(self.parents_tracker)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.vt = VariableBuilder(tx, self.source)(self.value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return type_dispatch(self, value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     unimplemented(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     raise Unsupported(msg)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code:
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]    File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return (z @ torch.einsum("i,j->ij", x, y).reshape(-1))
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting:
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     import torch._dynamo
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     torch._dynamo.config.suppress_errors = True
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ========== The above exception occurred while processing the following code ==========
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in <module>
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     main()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in main
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in print_performance
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in <listcomp>
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = model(*example_inputs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 29, in <lambda>
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return (z @ torch.einsum("i,j->ij", x, y).reshape(-1))
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] ==========
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] Traceback (most recent call last):
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 786, in _convert_frame
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     result = inner_convert(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return _compile(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return func(*args, **kwds)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     guarded_code = compile_inner(code, one_graph, hooks, transform)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     r = func(*args, **kwargs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     out_code = transform_code_object(code, transform)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     transformations(instructions, code_options)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return fn(*args, **kwargs)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     tracer.run()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     super().run()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     and self.step()
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     getattr(self, inst.opname)(inst)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 242, in impl
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.push(fn_var.call_function(self, self.popn(nargs), {}))
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in call_function
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     args = [v.realize() for v in args]
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 535, in <listcomp>
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     args = [v.realize() for v in args]
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self._cache.realize(self.parents_tracker)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     self.vt = VariableBuilder(tx, self.source)(self.value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     vt = self._wrap(value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return type_dispatch(self, value)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     unimplemented(
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     raise Unsupported(msg)
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] from user code:
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]    File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 10, in func_with_sparse
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     return (z @ torch.einsum("i,j->ij", x, y).reshape(-1))
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] You can suppress this exception and fall back to eager by setting:
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     import torch._dynamo
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822]     torch._dynamo.config.suppress_errors = True
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convert_frame.py:822] 
0.000150
I0713 17:21:46.286000 140632849619840 torch/_dynamo/logging.py:55] [0/0] Step 1: torchdynamo start tracing inner /global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py:34
Traceback (most recent call last):
  File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 38, in <module>
    main()
  File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 35, in main
    print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in print_performance
    timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 379, in <listcomp>
    timings = torch.tensor([timed(fn, args, times, device) for _ in range(repeat)])
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_inductor/utils.py", line 368, in timed
    result = model(*example_inputs)
  File "/global/u1/a/aglover/e3nn_benchmarking/minimal_sparse_compilation_error.py", line 35, in <lambda>
    print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
    return fn(*args, **kwargs)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 921, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
    return _compile(
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 676, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 262, in time_wrapper
    r = func(*args, **kwargs)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 535, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1036, in transform_code_object
    transformations(instructions, code_options)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 165, in _fn
    return fn(*args, **kwargs)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 500, in transform
    tracer.run()
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2149, in run
    super().run()
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 810, in run
    and self.step()
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 773, in step
    getattr(self, inst.opname)(inst)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 489, in wrapper
    return inner_fn(self, inst)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1238, in CALL_FUNCTION_EX
    isinstance(fn, GetAttrVariable)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 135, in __instancecheck__
    instance = instance.realize()
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 58, in realize
    self._cache.realize(self.parents_tracker)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/lazy.py", line 24, in realize
    self.vt = VariableBuilder(tx, self.source)(self.value)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
    vt = self._wrap(value)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 526, in _wrap
    keywords[k] = VariableBuilder(
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 269, in __call__
    vt = self._wrap(value)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 402, in _wrap
    return type_dispatch(self, value)
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 1069, in wrap_tensor
    unimplemented(
  File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 190, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: torch.compile does not support sparse Tensor with torch.sparse_csr layout

from user code:
   File "/global/homes/a/aglover/.conda/envs/e3nn_env/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 36, in inner
    return fn(*args, **kwargs)


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] TorchDynamo compilation metrics:
I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] Function, Runtimes (s)
I0713 17:21:46.396000 140632849619840 torch/_dynamo/utils.py:320] _compile.<locals>.compile_inner, 0.0000

### Minified repro

import torch
from torch._inductor.utils import print_performance
from functools import partial
import logging

def func_with_sparse(x : torch.tensor , y : torch.tensor , z : torch.tensor):
    """
    This peforms a contraction with a CSR matrix
    """
    return (z @ torch.einsum("i,j->ij", x, y).reshape(-1))

def main():
    torch._dynamo.config.verbose=True
    torch._logging.set_logs(dynamo = logging.INFO)

    device = "cuda"

    ### Tensors, 2 dense, 1 sparse
    x = torch.rand( 3,  dtype=torch.float32, device=device)
    y = torch.rand( 3,  dtype=torch.float32, device=device)
    z = torch.randn(size=(3,9), dtype=torch.float32, device=device).to_sparse_csr()

    ### Works with sparse arg 
    func_with_sparse_arg = partial(func_with_sparse, z=z)
    print_performance(lambda: func_with_sparse_arg(x, y), times=10, repeat=10)
    
    ### Suceeds compilation (fullgraph=False) with sparse arg
    compiled_fullgraph_false = torch.compile(func_with_sparse_arg, fullgraph=False, backend='eager') 
    print_performance(lambda: compiled_fullgraph_false(x, y), times=10, repeat=10)

    torch._dynamo.reset()

    ### Fails compilation (fullgraph=True) with sparse arg
    compiled_fullgraph_true = torch.compile(func_with_sparse_arg, fullgraph=True, backend='eager') 
    print_performance(lambda: compiled_fullgraph_true(x, y), times=10, repeat=10)

if __name__ == "__main__":
    main()


### Versions

(e3nn_env) aglover@login36:~/e3nn_benchmarking> python3 collect_env.py
Collecting environment information...
PyTorch version: 2.3.1.post300
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

OS: SUSE Linux Enterprise Server 15 SP4 (x86_64)
GCC version: (SUSE Linux) 12.3.0
Clang version: Could not collect
CMake version: version 3.20.4
Libc version: glibc-2.31

Python version: 3.10.14 | packaged by conda-forge | (main, Mar 20 2024, 12:45:18) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-5.14.21-150400.24.81_12.0.87-cray_shasta_c-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.5.82
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-PCIE-40GB
Nvidia driver version: 525.105.17
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:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             256
On-line CPU(s) list:                0-255
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 7713 64-Core Processor
CPU family:                         25
Model:                              1
Thread(s) per core:                 2
Core(s) per socket:                 64
Socket(s):                          2
Stepping:                           1
Frequency boost:                    enabled
CPU max MHz:                        2000.0000
CPU min MHz:                        1500.0000
BogoMIPS:                           3992.60
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
Virtualization:                     AMD-V
L1d cache:                          4 MiB (128 instances)
L1i cache:                          4 MiB (128 instances)
L2 cache:                           64 MiB (128 instances)
L3 cache:                           512 MiB (16 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-63,128-191
NUMA node1 CPU(s):                  64-127,192-255
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:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
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; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.3.1.post300
[conda] libmagma                  2.7.2                h173bb3b_2    conda-forge
[conda] libmagma_sparse           2.7.2                h173bb3b_3    conda-forge
[conda] libtorch                  2.3.1           cuda120_h2b0da52_300    conda-forge
[conda] mkl                       2023.2.0         h84fe81f_50496    conda-forge
[conda] numpy                     1.26.4          py310hb13e2d6_0    conda-forge
[conda] pytorch                   2.3.1           cuda120_py310h2c91c31_300    conda-forge

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip @ezyang @anijain2305 @chauhang @penguinwu