pybaum.registry
#
Module Contents#
Functions#
|
Create a pytree registry. |
- pybaum.registry.get_registry(types=None, include_defaults=True)[source]#
Create a pytree registry.
- Parameters
types (list) – A list strings with the names of types that should be included in the registry, i.e. considered containers and not leaves by the functions that work with pytrees. Currently we support: - “tuple” - “dict” - “list” -
collections.namedtuple
ortyping.NamedTuple
-None
-collections.OrderedDict
- “numpy.ndarray” - “jax.numpy.ndarray” - “pandas.Series” - “pandas.DataFrame”include_defaults (bool) – Whether the default pytree containers “tuple”, “dict” “list”, “None”, “namedtuple” and “OrderedDict” should be included even if not specified in types.
- Returns
A pytree registry.
- Return type