traveler.trees.AttrPyTree#

class traveler.trees.AttrPyTree(*args, **kwargs)[source]#

Bases: object

A JaxTree that allows for named attributes.

Allocate a pytree with empty data and metadata mappings.

Parameters:

  • *args (Any) – Positional initialization arguments accepted by the subclass.

  • **kwargs (Any) – Keyword initialization arguments accepted by the subclass.

Returns:

Newly allocated, empty instance of cls.

Return type:

AttrPyTree