Open asglover opened 1 month ago
I think this is probably something that your team is aware of, but I was asked to file a ticket. If there are any additional experiments or things you would like me to check, please let me know.
cc @pearu We have very limited support for torch.compile + sparse tensors. I don't think anybody is working on it at the moment.
Actually, @aartbik has been working on it. But don't expect speedups even if we can capture it https://github.com/pytorch/pytorch/pull/117907
Hi @ezyang Thanks for the info!
Hi @aartbik thank you for doing this work on the sparse compilers, it is much appreciated.
If you happen to be looking for real world applications / test cases for your compiler, I can share our application. We are going to solve it by writing a custom CUDA kernel, but the kernel could absolutely be generated by a sufficiently powerful sparse compiler.
Again, thank you for your work on this infrastructure!
@asglover I am very interested in getting more real world applications / test cases for our compiler, which is ramping up (see https://github.com/MPACT-ORG/mpact-compiler). Would you mind contacting me by email to make that work?
I'll write up a little summary of our application and the ways a compiler could help. Right off the bat, I would like to introduce @mitkotak who is looking at this problem specifically from a compiler perspective, and is very knowledgeable about compiler specifics. I'm focusing on the characteristics of the sparsity and trying to work out and implement by hand what I think should be performant. I'll email you when I get the time to put something together.
Thank you @asglover and @mitkotak. Looking forward to working together on sparse!
🐛 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
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
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/convertframe.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
I0713 17:21:46.240000 140632849619840 torch/_dynamo/convertframe.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
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
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
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
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/convertframe.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
I0713 17:21:46.260000 140632849619840 torch/_dynamo/convertframe.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
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
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
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 printperformance
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
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
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..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)
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