How to create custom trainable datatypes?

[NKlug]

I have a custom datatype consisting of a list of arrays. I want to use this datatype in a custom nnx.Module. During optimization, the internal parameters of this datatype should be adjusted. However, I have a hard time getting my module to "register" the datatype's internal parameters to its trainable parameters/state.
Here is some code:

class MyDType(object):
        def __init__(self, x: List[jnp.ndarray]):
                # x = preprocess(x)
                self.cores: List[jnp.ndarray] = [nnx.Param(y) for y in x]
        
        # some custom functions like __add__, __mul__, copy etc.


class MyModule(nnx.Module):
    def __init__(self, a: MyDType):
        self.order = len(a.cores)
        self.a = a

    def __call__(self, x):
        # some functionality, e.g.:
        assert x.shape[0] == self.order
        result = 0
        for i in range(self.order):
            result += jnp.sum(self.a.cores[i] ** x[i])
        return result


a = MyDType([jnp.ones(1), 2*jnp.ones(2), 3*jnp.ones(3)])
module = MyModule(a)

vars = nnx.state(module)

Here, vars in the last line is an empty state, while I would like it to contain the cores of a as parameters.

A workaround that I found is having MyDType inherit from nnx.Module. However, this does not feel like the proper way to do this, especially because MyDType represents a data type, not a module, and hence cannot be called. I didn't find anything on this in the docs.

How do I get the cores of MyDType show up in MyModule's state?

[cgarciae]

Hi @NKlug , you can inherit from nnx.Object which is the base for Module, Optimizer and other NNX types.

[NKlug]

Thanks, that works!