def METHOD_NAME(self) -> Tuple[list, dict]:
    """How to flatten a :class:`.JaxObject` instance into a pytree."""
    children = []
    aux_data = self.dict()
    for field_name in self.get_jax_field_names():
        field = getattr(self, field_name)
        sub_children, sub_aux_data = jax_tree_flatten(field)
        children.append(sub_children)
        aux_data[field_name] = sub_aux_data
    return children, aux_data