Which arguments should I guard in __init__ in pytrees?

[Jacob-Stevens-Haas]

The docs on custom pytrees warns that "the __init__ and __new__ methods of custom PyTree classes should generally avoid doing any array conversion or other input validation, or else anticipate and handle these special cases."

The two solutions involve

1. having tree_unflatten shadow the same logic as __init__, but using object.__new__() and assuming input validation has been done, or
2. special-casing when __init__ arguments will be of type object, Self, or None, then sidestep any TypeError in the logic, knowing that this isn't a "real" instance.

Which arguments need to allow object, Self, or None? All of them, just the first one, just the children, or just the aux_data?

[jakevdp]

Which arguments need to allow object, Self, or None? All of them, just the first one, just the children, or just the aux_data?

Just children

[jakevdp]

That said, I might take that suggestion out of the docs, because avoiding __init__ is a much better approach.

[Jacob-Stevens-Haas]

Thanks! That's clear. Though it makes me wonder - what's stopping jax from taking that advice? E.g. using the pattern suggested by @adonath for a dedicated hook like __mock_tree__ or __chart_children__ rather than relying on __init__ for object, None, or Self?

[adonath]

Just to extend on @jakevdp's answer: to avoid the issue with __init__ you can generally work with alternative constructors defined by @classmethod, as I suggested here: #25760 (comment). This pattern works really well in Jax. All input validation is done once explicitly by the alternative constructor. Once validated, the Jax serialization just goes through the "stupid" __init__ bypassing all validation.

[Jacob-Stevens-Haas]

Thanks, I should use this pattern more!