plum-parametric+jax-pytree.py

from jaxlib.xla_extension import DeviceArray as JaxArray
from jax.tree_util import register_pytree_node
from plum import dispatch, parametric
from typing import Tuple

import jax.numpy as jnp


def register_pytree_parametric(cls):
    register_pytree_node(
        cls,
        lambda xs: (tuple(xs.__dict__.values()), None),
        lambda _, xs: cls(xs)
    )
    return cls


@parametric
class Test:
    a: JaxArray
    
    @dispatch
    def __init__(self, args: Tuple):
        self.a, = args
    
    @dispatch
    def __init__(self, a):
        self.__init__((jnp.asarray(a),))