alpenstock.pytree.core

Classes:

Name	Description
AttrKey
FlattenFunc
ItemKey
LeafType
PathKey
Registration
TreeDef
TreePath
UnflattenFunc

Functions:

Name	Description
is_leaf
register_pytree_node	Register a custom pytree node type.
tree_flatten_with_path
tree_unflatten
unregister_pytree_node	Unregister a custom pytree node type.

Attributes:

Name	Type	Description
ChildrenAuxKey
PYTREE_REGISTRATIONS	dict[Type[Any], Registration]

ChildrenAuxKey module-attribute

ChildrenAuxKey = tuple[tuple[Any, ...], Any, tuple[PathKey, ...]]

PYTREE_REGISTRATIONS module-attribute

PYTREE_REGISTRATIONS: dict[Type[Any], Registration] = {}

AttrKey

Bases: PathKey

              flowchart TD
              alpenstock.pytree.core.AttrKey[AttrKey]
              alpenstock.pytree.core.PathKey[PathKey]

                              alpenstock.pytree.core.PathKey --> alpenstock.pytree.core.AttrKey
                


              click alpenstock.pytree.core.AttrKey href "" "alpenstock.pytree.core.AttrKey"
              click alpenstock.pytree.core.PathKey href "" "alpenstock.pytree.core.PathKey"
            

Methods:

Name	Description
__str__
get

Attributes:

Name	Type	Description
key	str

key class-attribute instance-attribute

key: str = field()

__str__

__str__() -> str

get

get(obj: Any) -> Any

FlattenFunc

Bases: Protocol

              flowchart TD
              alpenstock.pytree.core.FlattenFunc[FlattenFunc]

              

              click alpenstock.pytree.core.FlattenFunc href "" "alpenstock.pytree.core.FlattenFunc"
            

Methods:

Name	Description
__call__	The function to flatten a pytree node.

__call__

__call__(obj: Any) -> ChildrenAuxKey

The function to flatten a pytree node.

Parameters:

Name	Type	Description	Default
obj	Any	the pytree node to flatten	required

Returns:

Name	Type	Description
ChildrenAuxKey	ChildrenAuxKey	A tuple of (children, aux, path_keys), where children (tuple[Any, ...]): A tuple of child nodes. aux (Any): auxiliary data needed to reconstruct the node. path_keys (tuple[PathKey, ...]): A tuple of PathKey objects representing the path to each child.

ItemKey

Bases: PathKey

              flowchart TD
              alpenstock.pytree.core.ItemKey[ItemKey]
              alpenstock.pytree.core.PathKey[PathKey]

                              alpenstock.pytree.core.PathKey --> alpenstock.pytree.core.ItemKey
                


              click alpenstock.pytree.core.ItemKey href "" "alpenstock.pytree.core.ItemKey"
              click alpenstock.pytree.core.PathKey href "" "alpenstock.pytree.core.PathKey"
            

Methods:

Name	Description
__str__
get

Attributes:

Name	Type	Description
key	Any

key class-attribute instance-attribute

key: Any = field()

__str__

__str__() -> str

get

get(obj: Any) -> Any

LeafType

PathKey

Bases: ABC

              flowchart TD
              alpenstock.pytree.core.PathKey[PathKey]

              

              click alpenstock.pytree.core.PathKey href "" "alpenstock.pytree.core.PathKey"
            

Methods:

Name	Description
get

get abstractmethod

get(obj: Any) -> Any

Registration

Attributes:

Name	Type	Description
flatten	FlattenFunc
tree_def	Type[TreeDef]
type	Type[Any]
unflatten	UnflattenFunc

flatten class-attribute instance-attribute

flatten: FlattenFunc = field()

tree_def class-attribute instance-attribute

tree_def: Type[TreeDef] = field()

type class-attribute instance-attribute

type: Type[Any] = field()

unflatten class-attribute instance-attribute

unflatten: UnflattenFunc = field()

TreeDef

Methods:

Name	Description
__str__	The generic string representation of a TreeDef.

Attributes:

Name	Type	Description
aux	Any
children	tuple[TreeDef, ...]
is_leaf	bool
is_none	bool
type	Type[Any]

aux class-attribute instance-attribute

aux: Any = field(default=None)

children class-attribute instance-attribute

children: tuple[TreeDef, ...] = field(factory=tuple)

is_leaf property

is_leaf: bool

is_none property

is_none: bool

type class-attribute instance-attribute

type: Type[Any] = field()

__str__

__str__() -> str

The generic string representation of a TreeDef.

Developers are encouraged to override this method for custom pytree nodes.

TreePath

Methods:

Name	Description
__getitem__
__str__
__truediv__
get

Attributes:

Name	Type	Description
paths	tuple[PathKey, ...]

paths class-attribute instance-attribute

paths: tuple[PathKey, ...] = field(factory=tuple)

__getitem__

__getitem__(key: int | slice) -> PathKey | TreePath

__str__

__str__() -> str

__truediv__

__truediv__(key: PathKey) -> TreePath

get

get(obj: Any) -> Any

UnflattenFunc

Bases: Protocol

              flowchart TD
              alpenstock.pytree.core.UnflattenFunc[UnflattenFunc]

              

              click alpenstock.pytree.core.UnflattenFunc href "" "alpenstock.pytree.core.UnflattenFunc"
            

Methods:

Name	Description
__call__	The function to unflatten a pytree node.

__call__

__call__(aux: Any, children: tuple[Any, ...]) -> Any

The function to unflatten a pytree node.

Parameters:

Name	Type	Description	Default
aux	Any	the auxiliary data needed to reconstruct the node.	required
children	tuple[Any, ...]	A tuple of child nodes.	required

Returns:

Name	Type	Description
Any	Any	The reconstructed pytree node.

is_leaf

is_leaf(node: Any, path: TreePath) -> bool

register_pytree_node

register_pytree_node(type: Type[Any], flatten: FlattenFunc, unflatten: UnflattenFunc, tree_def: Type[TreeDef] = TreeDef) -> None

Register a custom pytree node type.

Parameters:

Name	Type	Description	Default
type	Type[Any]	The custom pytree node type to register.	required
flatten	FlattenFunc	The function to flatten the custom pytree node.	required
unflatten	UnflattenFunc	The function to unflatten the custom pytree node.	required
tree_def	Type[TreeDef]	The TreeDef class for the custom pytree node. Defaults to TreeDef.	TreeDef

tree_flatten_with_path

tree_flatten_with_path(tree: Any) -> tuple[list[TreePath], list[Any], TreeDef]

tree_unflatten

tree_unflatten(leaves: list[Any], tree_def: TreeDef) -> Any

unregister_pytree_node

unregister_pytree_node(type: Type[Any]) -> None

Unregister a custom pytree node type.

Parameters:

Name	Type	Description	Default
type	Type[Any]	The custom pytree node type to unregister.	required