pybaum.registry#

Module Contents#

Functions#

get_registry([types, include_defaults])

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 or typing.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

dict