is_pytree_dataclass_type

penzai.core.struct.is_pytree_dataclass_type(cls: type[Any]) → bool

Checks if a class was wrapped in the pytree_dataclass decorator.

Note that inheriting from a PyTree dataclass type does NOT produce another PyTree dataclass type, because JAX PyTree traversals depend on the specific class and not on its base classes.

Parameters:

cls – Class to check.

Returns:

True if this specific class was registred with @pytree_dataclass.