🏰 Backend API

sepes.backend_context(backend_name)

Context manager for switching the tree backend within a context.

Parameters:

backend_name (Literal['optree', 'jax']) – The name of the backend to switch to. available backends are optree and jax.

Example

Registering a custom tree class with optree backend:

>>> import sepes as sp
>>> import optree
>>> with sp.backend_context("optree"):
...     class Tree(sp.TreeClass):
...         def __init__(self, a, b):
...             self.a = a
...             self.b = b
...     tree = Tree(1, 2)
>>> optree.tree_flatten(tree, namespace="sepes")
([1, 2], PyTreeSpec(CustomTreeNode(Tree[('a', 'b')], [*, *]), namespace='sepes'))