JIT Error while defining inherited nn.Module class methods

[moehmeni]

Based on jax-ml/jax#10598 (comment) I am trying to make a class jit compatible by creating a pytree node, but I get:

TypeError: Cannot interpret value of type <class '__main__.TestModel'> as an abstract array; it does not have a dtype attribute

The error line is in the fit function when calling self.step.
Is there anything wrong with my implementation?

@register_pytree_node_class
@dataclass
class Parent(nn.Module):
    key: jax.random.PRNGKey
    params: dict = None

    @jax.jit
    def step(self, loss_fn, optimizer, opt_state, x, y):
       # ...

    @jax.jit
    def predict(self, x):
        return self.apply(self.params, x)

    def fit( # ... ):
            # ...
            for x, y in utils.data_loader(X, Y, batch_size):
                self.params, opt_state, loss_value = self.step(
                    loss, optimizer, opt_state, x, y
                ) # I get the error on this line
            # ...

    def tree_flatten(self):
        return (self.params,), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children, aux_data)

class TestModel(Parent):
 d_hidden: int = 64
 d_out: int = 1

 @nn.compact
 def __call__(self, x):
     # ...

model = TestModel(key=jax.random.PRNGKey(0))
model.fit(# ... )

[chiamp]

Is it possible to provide a simplified reproducible example? It's hard to debug since I can't execute this current example locally and reproduce the same error.