Can I use register_dataclass_node to pytree register classes generated by custom metaclass similar to dataclasses

[acecchini]

Hi,

I have constructed a metaclass which construct classes generating dynamically their __init__ method based on attributes defined in the class body.

Now I think it shouldn't be an issue to register these classes through a register_dataclass but the only thing that restrain me to do so is that the arguments passed to the __init__ method needs to be the one saved as instance attributes. That means for example I cannot register a class which store an instance attribute that reesults from some transformation of some of the __init__ parameters, i.e. when initializing tensor's weights, I know the shape of the tensor and I randomly sample from a certain probability distribution the actual weights before storing the tensor as an instance attribute.

I have looked up the register_dataclass function defintion and it is based on the register_dataclass_node method which is defined through a c++ api. Is it possible to use register_dataclass_node to register the classes I have described above?

Thanks

[jakevdp]

The only requirement for register_dataclass is that all relevant attributes be accepted as keyword arguments to the constructor, and that this is sufficient to reconstruct the instance of the class. If you have a class that meets these requirements, you should be fine using register_dataclass. If your class does not meet these requirements, you should use a different registration mechanism.