Repository: pytorch
Issue: torch.export does not support Sparse CSR tensors
Labels: onnx-needs-info,oncall: export

### 🐛 Describe the bug

Sparse CSR has experimental support as shown by the warning message `/opt/pytorch/torch/_dynamo/utils.py:1832: 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 /opt/pytorch/aten/src/ATen/SparseCsrTensorImpl.cpp:53.)` in the repro below.

The repro finds a scenario in which the current implementation fails to export it

```python
import torch

class MySParseThingy(torch.nn.Module):
    def forward(self, x):
        return [xi.to_sparse_csr() for xi in x]

x = [torch.randn(3, 3) for _ in range(3)]
ep = torch.export.export(MySParseThingy(), args=(x,))
```

Error:

```bash
/opt/pytorch/torch/_dynamo/utils.py:1832: 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 /opt/pytorch/aten/src/ATen/SparseCsrTensorImpl.cpp:53.)
  return getattr(args[0], node.target)(*args[1:], **kwargs)
Traceback (most recent call last):
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1399, in _dispatch_impl
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_ops.py", line 600, in __call__
    return self_._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: aten::_to_sparse_csr: attempted to run this operator with Meta tensors, but there was no abstract impl or Meta kernel registered. You may have run into this message while using an operator with PT2 compilation APIs (torch.compile/torch.export); in order to use this operator with those APIs you'll need to add an abstract impl.Please see the following doc for next steps: https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/pytorch/torch/_dynamo/utils.py", line 1832, in run_node
    return getattr(args[0], node.target)(*args[1:], **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 837, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1182, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 907, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1401, in _dispatch_impl
    return maybe_run_unsafe_fallback(not_implemented_error)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1387, in maybe_run_unsafe_fallback
    return run_fallback_kernel(self, func, flat_args, args_spec, error)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1666, in run_fallback_kernel
    return pytree.tree_map(map_out, r)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/utils/_pytree.py", line 900, in tree_map
    return treespec.unflatten(map(func, *flat_args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/utils/_pytree.py", line 736, in unflatten
    leaves = list(leaves)
             ^^^^^^^^^^^^
  File "/opt/pytorch/torch/_subclasses/fake_tensor.py", line 1654, in map_out
    and e._typed_storage()._cdata in storages
        ^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_tensor.py", line 242, in _typed_storage
    untyped_storage = self.untyped_storage()
                      ^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Cannot access storage of SparseCsrTensorImpl

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/opt/pytorch/repro_sparse_csr.py", line 8, in <module>
    ep = torch.export.export(MySParseThingy(), args=(x,))
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/__init__.py", line 174, in export
    return _export(
           ^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 705, in wrapper
    raise e
  File "/opt/pytorch/torch/export/_trace.py", line 688, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/exported_program.py", line 84, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 933, in _export
    gm_torch_level = _export_to_torch_ir(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 373, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
                        ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 1238, in inner
    result_traced = opt_f(*args, **kwargs)
                    ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 410, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 939, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 400, in _convert_frame_assert
    return _compile(
           ^^^^^^^^^
  File "/opt/conda/envs/ptca/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 686, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 265, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 541, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/bytecode_transformation.py", line 1078, in transform_code_object
    transformations(instructions, code_options)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 165, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 503, in transform
    tracer.run()
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 2202, in run
    super().run()
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 843, in run
    while self.step():
          ^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 757, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 481, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 1810, in CALL
    self.call_function(fn, args, kwargs)
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 697, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 703, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 2362, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 2476, in inline_call_
    tracer.run()
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 843, in run
    while self.step():
          ^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 757, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 481, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 1810, in CALL
    self.call_function(fn, args, kwargs)
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 697, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/variables/misc.py", line 586, in call_function
    return self.obj.call_method(tx, self.name, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/variables/tensor.py", line 461, in call_method
    return wrap_fx_proxy(
           ^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/variables/builder.py", line 1343, in wrap_fx_proxy
    return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/variables/builder.py", line 1428, in wrap_fx_proxy_cls
    example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 1712, in get_fake_value
    ret_val = wrap_fake_exception(
              ^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 1227, in wrap_fake_exception
    return fn()
           ^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 1713, in <lambda>
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 1846, in run_node
    unimplemented(make_error_message(e), from_exc=e)
  File "/opt/pytorch/torch/_dynamo/exc.py", line 203, in unimplemented
    raise Unsupported(msg) from from_exc
torch._dynamo.exc.Unsupported: Failed running call_method to_sparse_csr(*(FakeTensor(..., size=(3, 3)),), **{}):
Cannot access storage of SparseCsrTensorImpl

from user code:
   File "/opt/pytorch/repro_sparse_csr.py", line 5, in forward
    return [xi.to_sparse_csr() for xi in x]
  File "/opt/pytorch/repro_sparse_csr.py", line 5, in <listcomp>
    return [xi.to_sparse_csr() for xi in x]

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

(ptca) root@e9b784a5c219:/opt/pytorch# 
```

### Versions

main

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4