jaxtyping numpy support

[pablovela5620]

I'm wondering if pyserde support jaxtyping and this is a bug, or if it is just unsupported. The following works fine

import numpy as np
from serde import serde
from serde.json import to_json, from_json
from jaxtyping import Float32

@serde
class NPFoo:
    u: np.ndarray

npfoo = NPFoo(
    np.zeros((21, 3)).astype(np.float32),
)

json = to_json(npfoo)
print(json)

where as if I use a jaxtyping type (for dtype and shape checking with beartype) I get an error

@serde
class NPFoo:
    u: Float32[np.ndarray, "21 3"]

npfoo = NPFoo(
    np.zeros((21, 3)).astype(np.float32),
)

json = to_json(npfoo)
print(json)

Traceback (most recent call last):
  File "/home/pablo/0Dev/personal/volume-capture/tools/test_serde.py", line 16, in <module>
    json = to_json(npfoo)
           ^^^^^^^^^^^^^^
  File "/home/pablo/0Dev/personal/volume-capture/.pixi/envs/default/lib/python3.11/site-packages/serde/json.py", line 75, in to_json
    to_dict(obj, c=cls, reuse_instances=reuse_instances, convert_sets=convert_sets), **opts
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pablo/0Dev/personal/volume-capture/.pixi/envs/default/lib/python3.11/site-packages/serde/se.py", line 471, in to_dict
    return to_obj(  # type: ignore
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pablo/0Dev/personal/volume-capture/.pixi/envs/default/lib/python3.11/site-packages/serde/se.py", line 385, in to_obj
    raise SerdeError(e) from None
serde.compat.SerdeError: Unsupported type: ndarray

[yukinarit]

Hi @pablovela5620! jaxtyping is not supported by pyserde. Sorry I am not a user of jaxtyping/numpy, I can not help. If you're interested in contributing, I am happy to assist you though.

[pablovela5620]

That makes sense @yukinarit ! Id be happy to help, I've already tried looking through the code some. If you could point me in the right direction I can try to add support. Since the underlying type is np.ndarray I'm guessing it shouldn't be that difficult

[yukinarit]

Sure @pablovela5620

As you may already know, serialize/deserialize code for numpy can be found in numpy.py

You may need to modify se.py and de.py. Those two source files are about codegen.

If you have any questions, feel free to ask anything. @kmsquire is the author of the numpy support. Would be nice if he could also help us.