Cannot load saved PTensor

[tabasavr]

What happened?

Found this when tried to load a circuit with patametrized gates. Saves fine but fails to load

What did you expect to happen?

No response

Minimal Complete Verifiable Example

import quimb as qu
import quimb.tensor as qtn
import numpy as np

def f(x):
    return np.eye(3)

if __name__ == '__main__':
    ptensor = qtn.PTensor(f, [0], ("a", "b"))
    qu.save_to_disk(ptensor, "ptensor")
    qu.load_from_disk("ptensor")

Relevant log output

Traceback (most recent call last):
  File "/home/rsbat/dev/masters-enc/bug.py", line 13, in <module>
    qu.load_from_disk("ptensor")
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/quimb/utils.py", line 237, in load_from_disk
    return joblib.load(fname, **load_opts)
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/joblib/numpy_pickle.py", line 587, in load
    obj = _unpickle(fobj, filename, mmap_mode)
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/joblib/numpy_pickle.py", line 506, in _unpickle
    obj = unpickler.load()
  File "/usr/lib/python3.9/pickle.py", line 1212, in load
    dispatch[key[0]](self)
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/joblib/numpy_pickle.py", line 331, in load_build
    Unpickler.load_build(self)
  File "/usr/lib/python3.9/pickle.py", line 1717, in load_build
    setstate(state)
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/quimb/tensor/tensor_core.py", line 2832, in __setstate__
    self._data, self._inds, tags, self._left_inds = state
  File "/home/rsbat/dev/masters-enc/venv/lib/python3.9/site-packages/quimb/tensor/tensor_core.py", line 8826, in _data
    raise TypeError(
TypeError: You can only directly update the data of a ``PTensor`` with another ``PArray``. You can chain another function with the ``.modify(apply=fn)`` method. Alternatively you can convert this ``PTensor to a normal ``Tensor`` with ``t.unparametrize()``

Anything else we need to know?

No response

Environment

Python 3.9.5 quimb 1.4.0

[jcmgray]

Thanks for the issue, this should be quite fixable, hopefully just by adding __setstate__ and __getstate__ to PTensor and maybe PArray, (which customize pickling behaviour). I can try and take a look unless you want to attempt a fix yourself!

[tabasavr]

I've looked into it, overriding __setstate__ and __getstate__ in PTensor indeed seems to fix the issue. However there is a very weird case when it can't save it after optimizing a circuit, here's an example. Honestly I have no idea what's happening, but it looks like when not all tensors are optimized something goes wrong. And it doesn't look like it's caused by the fix. Don't really need this to work anymore, but let me know if you have any thoughts on this.

[jcmgray]

Hmm, I'll try and take a look. Looks like something weird to do with capturing locally defined functions somehow. Thanks for the PR fixing the original issue!