Cannot serialize NJTs using torch.save

[vmoens]

🐛 Describe the bug

torch.save fails when saving NJTs:

layout = torch.jagged
g0 = torch.zeros(1)
g1 = torch.zeros(2)
g = torch.nested.nested_tensor([g0, g1], layout=layout)
torch.save(g, "file.p")

which raises

  File "<ipython-input-8-962d8302f59b>", line 1, in <module>
    torch.save(g, "file.p")
  File "/Users/vmoens/venv/rl/lib/python3.11/site-packages/torch/serialization.py", line 726, in save
    _save(
  File "/Users/vmoens/venv/rl/lib/python3.11/site-packages/torch/serialization.py", line 954, in _save
    pickler.dump(obj)
TypeError: cannot pickle 'PyCapsule' object

cc @mruberry @mikaylagawarecki @cpuhrsch @jbschlosser @bhosmer @drisspg @soulitzer

Versions

PT nightly

[malfet]

@vmoens do you think it's a regression? Or it simply never worked?

[jbschlosser]

@vmoens I can't repro this on a fresh build; serialization works fine for me.

Maybe there's an issue with your version of dill? It looks like that error was fixed in dill 0.3.6 (PR).

[vmoens]

Can't repro either with a fresh env - closing

[davidberard98]

I have a repro now:

import torch
from torch.nested._internal.nested_tensor import nested_view_from_values_offsets, _nt_view_dummy

def fn(values, offsets):
    return nested_view_from_values_offsets(values, offsets)

fn_t = torch.fx.symbolic_trace(fn)

# this is needed to call sym_sizes_capsule, presumably:
# https://github.com/pytorch/pytorch/blob/d95a019704f526ee985cfe8d68261d34ddaf0e9d/torch/csrc/PyInterpreter.cpp#L812-L813
_nt_view_dummy().size()

torch.save(fn_t, "/tmp/py_capsule_repro.pt")

[jbschlosser]

Nice find @davidberard98! And to expand the original repro:

import torch

layout = torch.jagged
g0 = torch.zeros(1)
g1 = torch.zeros(2)
g = torch.nested.nested_tensor([g0, g1], layout=layout)
# accessing size caches a PyCapsule on the object
g.shape
# TypeError: cannot pickle 'PyCapsule' object
torch.save(g, "file.p")

[jbschlosser]

This is a general issue relating to subclasses that implement custom python sizes / strides. Opened #130154 to track this.