Default `pytree_dataclass` for inheritors of a common baseclass about penzai HOT 2 CLOSED

To anyone reading this, I solved this problem by defining my own metaclass:

metaclass_of_struct = type(pz.Struct)


class PytreeMeta(metaclass_of_struct):
    def __new__(cls, name, bases, namespace):
        if pz.Struct not in bases and not any(
            issubclass(base, pz.Struct) for base in bases
        ):
            bases = (pz.Struct,) + bases
        new_cls = type.__new__(cls, name, bases, namespace)
        return pz.pytree_dataclass(new_cls)

which my "overwrite" class for Struct declares as its metaclass.

from penzai.

Thanks for the question! As the error message suggests, this was actually an explicit design choice in Penzai to not automatically use @pz.pytree_dataclass for all subclasses. The main reasons:

• If you have multiple levels of Python object inheritance, it can be non-obvious which ones are dataclasses and which ones aren't. For instance, if you subclass an ordinary Python dataclass, you don't get a dataclass back. Requiring the explicit decorator makes it clear which classes are dataclasses and which aren't.
• Inheritance between dataclasses is really tricky; it can silently reorder the attributes that you've declared, or overwrite methods in unexpected ways. The @pz.pytree_dataclass decorator takes extra arguments that make sure you (and readers of your code) can tell what the decorator is going to do.

I'd advise against directly subclassing the metaclass of Struct in this way, and I can't guarantee that this will always work. Instead I'd recommend just adding the one-line @pz.pytree_dataclass decorator for each of the classes you want to define. Or, if you really want automatic inheritance of "pytree-dataclassness", you can just write your own metaclass from scratch, and not inherit from Struct at all. (Or you could use something like eqx.Module, which has implicit inheritance semantics.)

from penzai.