inheritence and @register_pytree_node_class

[jecampagne]

Hello,
It may be a stupid question raised by my ignorance, so do not worry too much.
Here is the use case

class Base():
    def __init__(self, *, gparams=None, **params):
        self._gparams = gparams   # Non-traced static parameters  (relevant for speed/accuracy trade-offs)
        self._params = params # Dictionary containing all traced parameters ( internal represention of this object)

    @property
    def params(self):
        return self._params

    @property
    def gparams(self):
        return self._gparams

   def func(self, x):
        return jnp.sin(x)

    def tree_flatten(self):
        # Define the children nodes of the PyTree that need tracing
        children = (self.params,)
        # Define auxiliary static data that doesn’t need to be traced
        aux_data = {"gparams": self.gparams}
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Recreates an instance of the class from flatten representation"""
        return cls(**(children[0]), **aux_data)

Then let us define

@register_pytree_node_class
class A(Base):
   def __init__(self, a=a, gparams=None):
        super().__init__(a=a,gparams=gparams)
   def do_stuff(self,x):
        return x*self.params["a"]

Now to make a test I need to make a A class that overload the f-function of Base.
I was trying so

register_pytree_node_class
class Amodif(A):
  def __init__(self, *args, **kwargs):
    super().__init__(*args,**kwargs)

  def f(self, x):
    raise NotImplementedError("%s does not implement f"%self.__class__.__name__)

But I get AttributeError: type object 'Amodif' has no attribute 'tree_unflatten'

So does someone knows how to inherit from A class and overlead Base::f ?

[jakevdp]

I ran your code and was unable to reproduce the error. I had to slightly modify it (a=a led to a NameError, so I changed it to a=None). I executed the following code and saw no error:

from jax.tree_util import tree_flatten
tree_flatten(Amodif())

To avoid the need for more guesswork, can you edit your question so that it contains a minimal reproducible example of the code you're running when you see the reported error? Thanks.

[jecampagne]

Ho. thanks. In fact, I realize that I certainly import a code that was not correctly setup (I tried to make a simplified version of a very long package ). Now it works.
I you think the question is a garbage you can close and erase it. I'm so sorry,

[jakevdp]

No worries, glad you figured out the issue!

[jecampagne]

thanks