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

quaternionic/quaternionic/arrays.py

Lines 17 to 30 in 36735fe

	def QuaternionicArray(jit=jit, dtype=float):
	"""Construct a quaternionic array type.
	This factory returns a `class` encapsulating a quaternionic array type,
	where the jit function and dtype are passed to this factory function, and
	used when creating the class. The returned class can then be used to
	instantiate actual arrays. This allows us to, for example, skip jit
	compilation and construct numpy arrays of dtype `object`, so that we can
	use more general python types than the standard numeric types — such as
	sympy expressions.
	"""
	class QArray(QuaternionPropertiesMixin(jit), QuaternionConvertersMixin(jit), np.ndarray):
	"""Subclass of numpy arrays interpreted as quaternions.

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