Hello, treepie nodes

%xmode Verbose

Exception reporting mode: Verbose

import jax

import treepie as tpy
from treepie import Aux, Node

Dataclasses

class Foo(tpy.Dataclass):
    x: Node[int]
    y: Aux[bool]

foo = Foo(1, True)

assert foo.x == 1
assert foo.y == True

print(foo)

Foo(x=1, y=True)

leaves, treedef = jax.tree.flatten(foo)

print(f"{leaves = }\n{treedef = }")

leaves = [1]
treedef = PyTreeDef(CustomNode(Foo[True], [*]))

print(jax.tree.unflatten(treedef, leaves))

Foo(x=1, y=True)

Frozen sequences

frozen_seq = tpy.FrozenSeq([2, True, None])

assert frozen_seq[0] == 2
assert frozen_seq[1] == True
assert frozen_seq[2] == None

print(frozen_seq)

FrozenSeq([2, True, None])

leaves, treedef = jax.tree.flatten(frozen_seq)

print(f"{leaves = }\n{treedef = }")

leaves = [2, True]
treedef = PyTreeDef(CustomNode(FrozenSeq[], [*, *, None]))

print(jax.tree.unflatten(treedef, leaves))

FrozenSeq([2, True, None])

Frozen mappings

frozen_map = tpy.FrozenMap({"a": 2, "b": True, "c": None})

assert frozen_map["a"] == 2
assert frozen_map["b"] == True
assert frozen_map["c"] == None

print(frozen_map)

FrozenMap({'a': 2, 'b': True, 'c': None})

leaves, treedef = jax.tree.flatten(frozen_map)

print(f"{leaves = }\n{treedef = }")

leaves = [2, True]
treedef = PyTreeDef(CustomNode(FrozenMap['a', 'b', 'c'], [*, *, None]))

print(jax.tree.unflatten(treedef, leaves))

FrozenMap({'a': 2, 'b': True, 'c': None})