π 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
, ordict
. A leaf is a non-container data structure like anint
,float
,string
, orArray
.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`()
withfield()
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 seeautoinit()
andfield()
.>>> 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 toget
,set
, orapply
a function to a leaf or a group of leaves usingleaf
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[0].get() Tree(a=1, b=None)
Note
AttributeError
is raised, If a method that mutates the instance is called directly. Instead usevalue_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 fromabc.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.**