pickling / multiprocessing issue

[balazon]

I was trying to replace numpy quaternion with quaternionic, but ran into an issue that my multiprocessing code can't run anymore, since quaternionic.array objects can't be pickled. Here's a short repro:

import quaternionic
from multiprocessing import Pool
import quaternion
import pickle


def test_pickle():
    data = {
        'q1': quaternionic.array.from_rotation_vector([0.1, 0.2, 0.3]),
    }
    with open('data.pickle', 'wb') as f:
        pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)

def test_pickle_np_quat():
    data = {
        'q1': quaternion.from_rotation_vector([0.1, 0.2, 0.3]),
    }
    with open('data.pickle', 'wb') as f:
        pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)

def process(q):
    return q.inverse

def test_pool():
    q1 = quaternionic.array.from_rotation_vector([0.1, 0.2, 0.3])
    q2 = quaternionic.array.from_rotation_vector([0.2, 0.3, 0.4])
    items = [q1, q2]

    with Pool(processes=1) as p:
        for res in p.imap_unordered(process, items, chunksize=1):
            print(res)

if __name__ == "__main__":
    test_pool() # does not work
    # test_pickle() # does not work
    # test_pickle_np_quat() # works

I tried to work around the problem using copyreg like this:

import copyreg
def pickle_quaternionic(obj):
    return quaternionic.array, (obj.ndarray.copy(),)
copyreg.pickle(quaternionic.array, pickle_quaternionic)

Since that would just use ndarray for the serialization, but that does not work either, in both cases I get this output:

Traceback (most recent call last):
  File "C:\work\quaternionic\test_quaternionic.py", line 41, in <module>
    test_pool() # does not work
    ~~~~~~~~~^^
  File "C:\work\quaternionic\test_quaternionic.py", line 37, in test_pool
    for res in p.imap_unordered(process, items, chunksize=1):
               ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\balazon\miniforge3\envs\py13\Lib\multiprocessing\pool.py", line 873, in next
    raise value
  File "C:\Users\balazon\miniforge3\envs\py13\Lib\multiprocessing\pool.py", line 540, in _handle_tasks
    put(task)
    ~~~^^^^^^
  File "C:\Users\balazon\miniforge3\envs\py13\Lib\multiprocessing\connection.py", line 206, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
                     ~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "C:\Users\balazon\miniforge3\envs\py13\Lib\multiprocessing\reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
    ~~~~~~~~~~~~~~~~~~~~~~~^^^^^
AttributeError: Can't get local object 'QuaternionicArray.<locals>.QArray'

The problem stems from QArray class being in QuaternionicArray function's local scope https://github.com/moble/quaternionic/blob/36735feb464a10aaf86538d0d27f98ca3422bd7f/quaternionic/arrays.py#L17-L30 For pickling to work, it should be globally importable.

Chatgpt recommends this: You could lift QArray to a top-level definition in arrays.py:

class QArray(...):
    ...

And have QuaternionicArray return it with modifications if needed. That would allow it to be pickleable

I really don't know if this is the best way and why QArray is currently inside QuaternionicArray function, I see it has a jit, and a dtype parameter. I trust you can solve this issue, but right now it seems I can't really make the change to quaternionic in this state.

Thanks in advance