📍 Module API#

class sepes.TreeClass(*a, **k)#

Convert a class to a pytree by inheriting from TreeClass.

A pytree is any nested structure of containers and leaves. A container is a pytree can be a container or a leaf. Container examples are: a tuple, list, or dict. A leaf is a non-container data structure like an int, float, string, or Array. TreeClass is a container pytree that holds other pytrees in its attributes.

Note

TreeClass is immutable by default. This means that setting or deleting attributes after initialization is not allowed. This behavior is intended to prevent accidental mutation of the tree. All tree modifications on TreeClass are out-of-place. This means that all tree modifications return a new instance of the tree with the modified values.

There are two ways to set or delete attributes after initialization:

Using at property to modify an existing leaf of the tree.

>>> import sepes as sp
>>> class Tree(sp.TreeClass):
...     def __init__(self, leaf: int):
...         self.leaf = leaf
>>> tree = Tree(leaf=1)
>>> new_tree = tree.at["leaf"].set(100)
>>> tree is new_tree  # new instance is created
False

Using value_and_tree() to call a method that mutates the tree. and apply the mutation on a copy of the tree. This option allows writing methods that mutate the tree instance but with these updates applied on a copy of the tree.

>>> import sepes as sp
>>> class Tree(sp.TreeClass):
...     def __init__(self, foo: int):
...         self.foo = foo
...     def add_bar(self, value:int) -> None:
...         # this method mutates the tree instance
...         # and will raise an `AttributeError` if called directly.
...         setattr(self, "bar", value)
>>> tree = Tree(foo=1)
>>> # now lets try to call `add_bar` directly
>>> tree.add_bar(value=100)  
Cannot set attribute value=100 to `key='bar'`  on an immutable instance of `Tree`.
>>> output, tree_ = sp.value_and_tree(lambda T: T.add_bar(100))(tree)
>>> tree, tree_
(Tree(foo=1), Tree(foo=1, bar=100))

This pattern is useful to write freely mutating methods, but with The expense of having to call through at[“method_name”] instead of calling the method directly.

Note

sepes offers two methods to construct the __init__ method:

Manual __init__ method

>>> import sepes as sp
>>> class Tree(sp.TreeClass):
...     def __init__(self, a:int, b:float):
...         self.a = a
...         self.b = b
>>> tree = Tree(a=1, b=2.0)

Auto generated __init__ method from type annotations.

Using autoinit() decorator where the type annotations are used to generate the __init__ method. autoinit`() with field() objects can be used to apply functions on the field values during initialization, support multiple argument kinds, and can apply functions on field values on getting the value. For more details see autoinit() and field().
```
>>> import sepes as sp
>>> @sp.autoinit
... class Tree(sp.TreeClass):
...     a:int
...     b:float
>>> tree = Tree(a=1, b=2.0)
```

Note

Leaf-wise math operations are supported using leafwise decorator. leafwise decorator adds __add__, __sub__, __mul__, … etc to registered pytrees. These methods apply math operations to each leaf of the tree. for example:

>>> @sp.leafwise
... class Tree(sp.TreeClass):
...     def __init__(self, a:int, b:float):
...         self.a = a
...         self.b = b
>>> tree = Tree(a=1, b=2.0)
>>> tree + 1  # will add 1 to each leaf
Tree(a=2, b=3.0)

Note

Advanced indexing is supported using at property. Indexing can be used to get, set, or apply a function to a leaf or a group of leaves using leaf name, index or by a boolean mask.

>>> class Tree(sp.TreeClass):
...     def __init__(self, a:int, b:float):
...         self.a = a
...         self.b = b
>>> tree = Tree(a=1, b=2.0)
>>> tree.at["a"].get()
Tree(a=1, b=None)
>>> tree.at["a"].get()
Tree(a=1, b=None)

Note

AttributeError is raised, If a method that mutates the instance is called directly. Instead use value_and_tree() to call the method on a copy of the tree. value_and_tree() calls the function on copied input arguments to ensure non-mutating behavior.

>>> import sepes as sp
>>> class Counter(sp.TreeClass):
...     def __init__(self, count: int):
...         self.count = count
...     def increment(self, value):
...         self.count += value
...         return self.count
>>> counter = Counter(0)
>>> sp.value_and_tree(lambda C: C.increment(1))(counter)
(1, Counter(count=1))

Note

TreeClass inherits from abc.ABC meaning that it can abc features like @abc.abstractmethod can be used to define abstract behavior that can be implemented by subclasses.

Warning

The structure should be organized as a tree. In essence, cyclic references are not allowed. The leaves of the tree are the values of the tree and the branches are the containers that hold the leaves.

property at: at[Self]#

Immutable out-of-place indexing.

.at[***].get():
Return a new instance with the value at the index otherwise None.
.at[***].set(value):
Set the value and return a new instance with the updated value.
.at[***].apply(func):
Apply a func and return a new instance with the updated value.

Acceptable indexing types are:

str for mapping keys or class attributes.
int for positional indexing for sequences.
... to select all leaves.
a boolean mask of the same structure as the tree
a tuple of the above types to index multiple keys at same level.

Example

>>> import sepes as sp
>>> class Tree(sp.TreeClass):
...    def __init__(self, a:int, b:float):
...        self.a = a
...        self.b = b
...    def add(self, x: int) -> int:
...        self.a += x
...        return self.a
>>> tree = Tree(a=1, b=2.0)
>>> tree.at["a"].get()
Tree(a=1, b=None)
>>> tree.at["a"].set(100)
Tree(a=100, b=2.0)
>>> tree.at["a"].apply(lambda x: 100)
Tree(a=100, b=2.0)

Note

pytree.at[*][**] is equivalent to selecting pytree.*.** .
pytree.at[*, **] is equivalent selecting pytree.* and pytree.**

📍 Module API

Contents

📍 Module API#