Composite nodes¶

Imports for this tutorial:

import math
from typing import Any, Protocol

import jax
import jax.numpy as jnp

import treepie as tpy
from treepie import Aux, Node

What are composite nodes?¶

< Draw comparison with properties >

A composite node is an instance of any python class that implements the following protocol:

class CompositeNode(Protocol):
    def __unwrap__(self, parent, /): ...

When a composite node is set as a child-node in a treepie dataclass, it has a dual nature. When accessed as an attribute, the __unwrap__ dunder method is called and its value is returned. On the other hand, under jax tree operations (jax.tree.map, jax.tree.flatten, etc.), the node retains its full composite structure.

Examples¶

Example 1¶

Let’s start with a simple example. We want to parameterize a positive float value as exp(log_value), using a composite node. (Note: The __repr__ method is not necessary; we add it here to help with the pedagogy.)

class PositiveFloat:
    log_value: float

    def __init__(self, log_value: float) -> None:
        self.log_value = log_value

    def __unwrap__(self, parent) -> float:
        return math.exp(self.log_value)

    def __repr__(self) -> str:
        return f"PositiveFloat(log_value={self.log_value!r})"

Now let’s assign a PositiveFloat object as a node in a treepie dataclass:

class Foo(tpy.Dataclass):
    x: Node[float]

    def x_sq(self) -> float:
        return self.x**2


log_x = 0.5
x_value = math.exp(log_x)
composite_x = PositiveFloat(log_value=log_x)

foo = Foo(x=composite_x)

When accessed as an attribute, the __unwrap__ method is called and returned, so foo.x equals composite_x.__unwrap__(foo) (which equals x_value).

assert x_value == foo.x
assert x_value == getattr(foo, "x")
assert x_value == object.__getattribute__(foo, "x")
assert x_value**2 == foo.x_sq()

However, when you flatten foo, you get composite_x as the node.

print(jax.tree.leaves(foo))

[PositiveFloat(log_value=0.5)]

Here composite_x is a pytree leaf, so it didn’t get flattened further. On the other hand, if composite_x itself had been an internal JAX node, it would’ve been recursively flattened down to its leaves.

Example 2¶

Next, let’s look at a slightly more complicated example, where (a) the composite node is an internal jax node, and (b) the __unwrap__ method uses the parent position-only parameter. Let’s define a model to implement the transformationx.T @ inp @ y:

class Bar(tpy.Dataclass):
    x: Node[jax.Array]
    y: Node[jax.Array]

    def __call__(self, inp: jax.Array) -> jax.Array:
        return self.x.T @ inp @ self.y

Example 2.1: Now, let’s say we want to set y = x, and we don’t want them to be independent parameters:

class LinkToParentDotX(tpy.Dataclass):
    def __unwrap__(self, parent):
        return parent.x


x = jnp.ones(shape=(1, 3))
bar = Bar(x=x, y=LinkToParentDotX())

assert bar.y is bar.x
assert jax.tree.leaves(bar) == [x]

What’s happening here

bar.y (and also self.y within Bar.__call__) evaluates to the output of the composite node’s __unwrap__ method. So, bar.y is simply bar.x.

The pytree structure of bar is as follows:

── root node (`Bar` object):
   ├── leaf labeled `x` (`jax.Array` object) 
   └── node labeled `y` (`LinkToParentDotX` object; doesn't contain any child nodes)

Since LinkToParentDotX() doesn’t have any child nodes, it flattens out to an empty list. Consequently, when bar is flattened, only one copy of x shows up among its leaves.

Example 2.2: Such parameter sharing can be done for a generic attribute name as follows:

class SymLink(tpy.Dataclass):
    orig_attr_name: Aux[str]

    def __unwrap__(self, parent):
        return getattr(parent, self.orig_attr_name)


x = jnp.ones(shape=(1, 3))
bar = Bar(x=x, y=SymLink("x"))
assert bar.y is bar.x
assert jax.tree.leaves(bar) == [x]

Example 2.3: Maybe we want to parameterize y as x + exp(t), where x and t are the model’s free parameters

class CompositeY(tpy.Dataclass):
    t: Node[jax.Array]

    def __unwrap__(self, parent):
        return parent.x + jnp.exp(self.t)


x = jnp.ones(shape=(1, 3))
t = jnp.zeros(shape=(1, 3))
bar = Bar(x=x, y=CompositeY(t=t))
assert jnp.all(bar.y == x + jnp.exp(t))
assert jax.tree.leaves(bar) == [x, t]

All this was possible without having to modify the definition of the class Bar or subclass it. No boilerplate to unwrap the composite node within Bar.__call__ or somewhere else (as in paramax).

Accessing a child node without unwrapping it¶

Sometimes, one may want to access a child of a treepie node, without wrapping it. This can be done in any of the following ways:

bar.__getchild__("y")

CompositeY(t=Array([[0., 0., 0.]], dtype=float32))

tpy.get_children_map(bar)

mappingproxy({'x': Array([[1., 1., 1.]], dtype=float32),
              'y': CompositeY(t=Array([[0., 0., 0.]], dtype=float32))})

tpy.get_children_map(bar)["y"]

CompositeY(t=Array([[0., 0., 0.]], dtype=float32))

bar >> "y"

CompositeY(t=Array([[0., 0., 0.]], dtype=float32))

Adding tags and properties¶

class NonTrainable[T: Any](tpy.Dataclass):
    value: Node[T]

    def __unwrap__(self, parent: Any) -> T:
        return self.value


class MaybeTrainable[T: Any](tpy.Dataclass):
    value: Node[T]
    is_trainable: Aux[bool]

    def __unwrap__(self, parent: Any) -> T:
        return self.value


class ArrayWithNamedAxes(tpy.Dataclass):
    value: Node[jax.Array]
    named_axes: Aux[tuple[tuple[int, str]]]

    def __unwrap__(self, parent: Any) -> jax.Array:
        return self.value

FrozenSeq and FrozenMap also unwrap composite nodes¶

class CompositeY_Alt(tpy.Dataclass):
    t: Node[jax.Array]

    def __unwrap__(self, parent):
        return parent[0] + jnp.exp(self.t)


x = jnp.ones(shape=(1, 3))
t = jnp.zeros(shape=(1, 3))

frozen_seq = tpy.FrozenSeq([x, CompositeY_Alt(t)])

assert jnp.all(frozen_seq[1] == frozen_seq[0] + jnp.exp(t))
assert jax.tree.leaves(frozen_seq) == [x, t]

class CompositeY_Alt2(tpy.Dataclass):
    t: Node[jax.Array]

    def __unwrap__(self, parent):
        return parent["x"] + jnp.exp(self.t)


x = jnp.ones(shape=(1, 3))
t = jnp.zeros(shape=(1, 3))

frozen_map = tpy.FrozenMap({"x": x, "y": CompositeY_Alt2(t)})

assert jnp.all(frozen_map["y"] == frozen_map["x"] + jnp.exp(t))
assert jax.tree.leaves(frozen_map) == [x, t]