stringjax_tools.pytrees.make_pytree_flatteners

stringjax_tools.pytrees.make_pytree_flatteners#

make_pytree_flatteners(myclass, *, static_keys=(), ignore_keys=(), ignore_defaults=None, static_types=(<class 'str'>, <class 'bool'>), validate_static=True)#

Build class-specific pytree flatten/unflatten callables.

This convenience wrapper is useful when a project wants explicit registration calls without first naming a shared PytreePolicy.

Return type:

tuple[Callable[[Any], tuple[tuple[Any, ...], tuple[tuple[str, ...], tuple[tuple[str, Any], ...]]]], Callable[[tuple[tuple[str, ...], tuple[tuple[str, Any], ...]], tuple[Any, ...]], Any]]