Ancestry aware nodes

NOTE: Throughout this tutorial, “ancestry” is used in the pytree sense and not in the python class-inheritance sense. if foo = {"bar": 1}, then foo is an ancestor of foo["bar"].

An aside before we begin…
We’ll use this simple utility function to check if two or more lists are element-wise identical. This is a stronger condition than equality of lists.

def elementwise_identical_lists(ref_list: list, /, *other_lists: list) -> bool:
    assert other_lists, 'At least one "other list" must be provided'

    for L in other_lists:
        if len(L) != len(ref_list):
            return False

    for ref_elem, *other_elems in zip(ref_list, *other_lists):
        for elem in other_elems:
            if elem is not ref_elem:
                return False

    return True

Imports for this tutorial:

import jax
import jax.numpy as jnp

import treepie as tpy
from treepie import Node

Quick overview

Treepie’s pytree nodes are ancestry-aware. When defining a treepie dataclass, labels can be added to pytree nodes of that type as follows:

class LabeledNode(tpy.Dataclass):
    __node_labels__ = ("label_1", "label_2")

    # Rest of the class definition.

Specification (Move this to later?)

The __node_labels__ attribute should be defined without any type annotations and must be set to either (i) None or (ii) a tuple of str objects. An error will be raised during class creation if __node_labels__ doesn’t satisfy these conditions.

When a pytree dataclass is subclassed, node labels will be inherited (in the standard python-inheritance sense). Setting __node_labels__ to None is a way to block such inheritance.

A labeled ancestor of a node can be accessed as follows:

tpy.get_labeled_ancestor(descendant_node, label_str)

This will return the closest ancestor (in the pytree sense) of descendant_node with label_str as one of its labels, provided that every node of the pytree between the ancestor and descendant_node (including the ancestor and the descendant themselves) are ancestry aware nodes.

A read-only map (from label to ancestor) of all the visible ancestors of a descendant can be retrieved as:

tpy.get_labeled_ancestors_map(descendant_node)

So, a labeled ancestor can also be accessed as:

tpy.get_labeled_ancestors_map(descendant_node)[label_str]

Example

class Foo(tpy.Dataclass):
    bar: Node


class Bar(tpy.Dataclass):
    __node_labels__ = ("custom_label",)
    baz: Node


class Baz(tpy.Dataclass):
    quux: Node


class Quux(tpy.Dataclass): ...


foo = Foo(Bar(Baz(Quux())))

Here, foo is a pytree with the following tree structure: Foo object -> Bar object -> Baz object -> Quux object. The Bar object has a label "custom_label". One can access the Bar object from its descendants:

assert tpy.get_labeled_ancestor(foo.bar.baz, "custom_label") is foo.bar
assert tpy.get_labeled_ancestor(foo.bar.baz.quux, "custom_label") is foo.bar

tpy.get_labeled_ancestors_map(foo.bar.baz.quux)

mappingproxy({'custom_label': Bar(baz=Baz(quux=Quux()))})

A quick use case: weight sharing

Let’s say we want to create a machine learning model where the same weight-parameter is to be shared by different parts of a pytree. This can be accomplished as follows:

class SharedWeightRef(tpy.Dataclass):
    # This is a composite node. It contains no leaves.
    # When accessed as an attribute in a treepie dataclass, it
    # looks up an ancestor with the label "shared_weight_vendor"
    # and returns the shared_weight attribute from that ancestor.

    def __unwrap__(self, parent):
        weight_vendor = tpy.get_labeled_ancestor(self, "shared_weight_vendor")

        return weight_vendor.shared_weight


class SampleModel(tpy.Dataclass):
    # Here were marking this dataclass with the "shared_weight_vendor"
    # label. So, ancestor lookup from descedant SharedWeightRef
    # instances will find the corresponding LanguageModel instance.

    __node_labels__ = ("shared_weight_vendor",)

    shared_weight: Node[jax.Array]

    # Rest of the class definition

SharedWeightRef() can be used wherever one needs the shared weight, within any subtree of SampleModel.

Let’s make this example slightly more concrete (closely following the setup of the example for shared layers in equinox’s documentation). Let’s say two different dataclasses called Embedder and Linear, which implement some computational logic, have already been defined.

class SomeEmbedderSubmodule(tpy.Dataclass):
    some_wgt_param: Node[jax.Array]  # <--- wgt param 1


class Embedder(tpy.Dataclass):
    some_embedder_submodule: Node[SomeEmbedderSubmodule]


class Linear(tpy.Dataclass):
    some_other_wgt_param: Node[jax.Array]  # <--- wgt param 2

There are two weight parameters explicitly declared above: some_wgt_param in SomeEmbedderSubmodule and some_other_wgt_param in Linear. There could be other parameters in these dataclasses. There could also be different methods in these models. But we leave those out in this minimal example.

Now, we want to create a new dataclass called LanguageModel, which contains an Embedder instance and a Linear instance as subtrees, such that the two weight parameters are shared. Such sharing can be accomplished as follows (leaving out other components and methods of LanguageModel in this minimal example):

class SharedWeightRef(tpy.Dataclass):
    def __unwrap__(self, parent):
        weight_vendor = tpy.get_labeled_ancestor(self, "shared_weight_vendor")

        return weight_vendor.shared_weight


class LanguageModel(tpy.Dataclass):
    __node_labels__ = ("shared_weight_vendor",)

    shared_weight: Node[jax.Array]

    embedder: Node[Embedder]
    linear: Node[Linear]

    def __init__(self, shared_weight: jax.Array) -> None:
        self.shared_weight = shared_weight

        self.embedder = Embedder(
            some_embedder_submodule=SomeEmbedderSubmodule(
                some_wgt_param=SharedWeightRef()  # <-- we're assigning a reference here
            )
        )

        self.linear = Linear(
            some_other_wgt_param=SharedWeightRef()  # <-- and here
        )

By setting the embedder and linear weights to SharedWeightRef objects, we have encoded the weight-sharing logic directly into the pytree structure of LanguageModel objects. Consequently, as demonstrated below, treepie nodes need no special processing to make this parameter sharing work as intended.

Differences from other JAX frameworks

  • The weight sharing logic is expressed explicitly in the definition of SharedWeightRef, as opposed to implicitly using references to the same python object (as in penzai or flax, both of which require special processing of their pytrees).

  • The weight sharing logic is fully-contained within the pytree creation process. One does not need to implement the logic for copying the weight to the right locations, after the appropriate JAX transform boundaries have been crossed. In other words, the weight sharing logic doesn’t need to be implemented within the computational logic (e.g. within __call__) as in equinox.

Note

In the example above, it was possibly to just pass SharedWeightRef objects as inputs for initializing the Embedder and Linear objects. In practice, Embedder and Linear might have custom init methods, which don’t accept weights as inputs. In such cases, one may have to perform tree surgery within __init__ to place the SharedWeightRef objects in the right places. But this logic still resides within __init__ as opposed to, say, __call__.

< talk about how treescope works as it does with penzai >

Before proceeding further, let’s make sure that the weight sharing is working as intended:

wgt = jnp.zeros(shape=(3, 3))

lang_mdl = LanguageModel(wgt)

# Ensure that the references are pointing to the shared weight array
assert (
    lang_mdl.embedder.some_embedder_submodule.some_wgt_param is lang_mdl.shared_weight
)
assert lang_mdl.linear.some_other_wgt_param is lang_mdl.shared_weight

# Ensure that the references don't count as leaves
assert jax.tree.leaves(lang_mdl.embedder) == []
assert jax.tree.leaves(lang_mdl.linear) == []

# Ensure that exactly one copy of the shared weight exists among the leaves of lang_mdl
assert elementwise_identical_lists(
    jax.tree.leaves(lang_mdl),
    [lang_mdl.shared_weight],
    [wgt],
)

# Pretty printing the path of the array-leaf:
print(
    "treepath of shared weight = "
    "lang_mdl" + jax.tree_util.keystr(jax.tree.leaves_with_path(lang_mdl)[0][0])
)

treepath of shared weight = lang_mdl.__getchild__('shared_weight')

This doesn’t yet demonstrate that the weight-sharing is encoded into pytree structure. What happens when you keep the tree structure and change the leaf? Will the weight-references in the resulting pytree return the new weight (as they should) or be stuck returning the old weight (as they shouldn’t)? Let’s find out:

orig_leaves, treedef = jax.tree.flatten(lang_mdl)

wgt_new = jnp.ones(shape=(2, 2))
assert wgt_new.shape != wgt.shape
lang_mdl_new = jax.tree.unflatten(treedef, [wgt_new])

assert (
    lang_mdl_new.embedder.some_embedder_submodule.some_wgt_param
    is lang_mdl_new.shared_weight
)
assert lang_mdl_new.linear.some_other_wgt_param is lang_mdl_new.shared_weight

assert jax.tree.leaves(lang_mdl_new.embedder) == []
assert jax.tree.leaves(lang_mdl_new.linear) == []
assert elementwise_identical_lists(
    jax.tree.leaves(lang_mdl_new),
    [lang_mdl_new.shared_weight],
    [wgt_new],
)

assert jax.tree.leaves(lang_mdl)[0] is not jax.tree.leaves(lang_mdl_new)[0]

It works as intended! The fact that this works under tree flattens and unflattens means that it’ll work under JAX-transforms as well. When lang_mdl crosses a JAX-transform boundary, JAX will (a) flatten lang_mdl, (b) replace the shared weight with an appropriate tracer, and (c) unflatten the tree with the tracer-leaf. Past the JAX-transform boundary, the weight references will point to the right tracer object. No special processing or boilerplate necessary.

Let’s test that claim. test_func is a simple function that just returns the sum of the two weights within embedder and linear. The output should be two times the shared weight. We setup this function for jit-compilation with the @jax.jit decorator and call it on a test input:

@jax.jit
def test_func(lang_mdl: LanguageModel):
    print("COMPILATION TRIGGERED")

    return (
        lang_mdl.embedder.some_embedder_submodule.some_wgt_param
        + lang_mdl.linear.some_other_wgt_param
    )


test_shared_wgt = jnp.array([[0, 1], [2, 3]])
test_input = LanguageModel(shared_weight=test_shared_wgt)

print("Yell if a compilation is triggered...", end=" ")
assert jnp.all(test_func(LanguageModel(test_shared_wgt)) == test_shared_wgt * 2)

Yell if a compilation is triggered... COMPILATION TRIGGERED

We see that this call triggers a jit-compilation and the result is as expected.

Next, let’s call this function on a different input, with the same dtype and shape. We want (a) jit compitation to not be triggered again and (b) the output to correspond to the new test input.

test_shared_wgt = jnp.array([[4, 5], [6, 8]])
test_input = LanguageModel(shared_weight=test_shared_wgt)

print("Yell if a compilation is triggered...", end=" ")
assert jnp.all(test_func(LanguageModel(test_shared_wgt)) == test_shared_wgt * 2)

Yell if a compilation is triggered...

Success!

Ancestry semantics and treepie’s implementation

“If the implementation is hard to explain, it’s a bad idea.” - Zen of Python

Ancestry awareness is one of the most powerful features offered treepie. It can help avoid threading random keys (just query an ancestor labeled "random_key_vendor", avoid threading state objects (same as before, except the label can be "state_manager"). It can enable communication across different parts of a pytree, greatly enhancing the practical expressivity of pytrees. (Sidenote: If you’re into design patterns and such, treepie’s ancestry awareness is a way to emulate “algebraic effects”, which is a modern concept in functional programming to declare and handle side-effects in a clean, safe, modular, and composable way).

Ancestry awareness is also the most unusual feature within the JAX-ecosystem (the fact that the previous example worked as intended might be surprising). Although the feature doesn’t have much of a learning curve, it has some potentially surprising (borderline magic) behavior. So, it may be worthwhile to read through the complete feature-specification and implementation (covered next).

Some questions one might have at this point in the tutorial are:

  • What will happen if I place the same node within two different pytrees and query its ancestor?

  • If a pytree is deep-copied, what’s the ancestry of the nodes in the copied tree?

  • Does the call tpy.get_labeled_ancestor(self, "shared_weight_vendor") within SharedWeightRef.__unwrap__ not constitute a side-effect? (Answer: Yes, it does)

  • Is that not unsafe? (Answer: No, it isn’t).

Let’s go over some specifications and then answer these questions.

Warning: What you set may not be what you get

Setup:

  1. Let child_candidate and parent both be treepie nodes (i.e., dataclass instances, or instances of treepie.FrozenMap or treepie.FrozenSeq). Let child_candidate be assigned as a child-node within parent.

    # This can happen, for example as

class ParentCls(tpy.Dataclass):
    attr_name: Node

parent = ParentCls(child_candidate)

# or as

parent = tpy.FrozenMap({key: child_candidate})

# or as

parent = tpy.FrozenSeq([child_candidate])

  2. Next, let the the child be accessed as parent.__getchild__("attr_name") or as parent.__getchild__(key) or as parent.__getchild__(0).

Result: The value returned by the access may not necessarily be child_candidate. It could be a shallow copy of child_candidate.

Note: Accesses of the form parent.attr_name, parent[key], and parent[0] first call __getchild__ to retrieve the child, before unwrapping the returned child if it is composite. So, the caveat about shallow copying applies there too. Likewise for accesses of the form tpy.get_children_map(parent)[<"attr_name"/key/0>], object.__getattribute__(parent, "attr_name"), parent.__getitem__(<key/0>), etc.

Definition: Visible ancestor

foo is a visible ancestor of bar with label label_str if and only if all these conditions are satisfied:

  • foo is an instance of a treepie dataclass with label_str as one of its node labels

  • bar is a treepie node

  • bar can be reached from foo, through an uninterrupted sequence of treepie nodes as follows:

    assert bar is foo.__getchild__(blah_1). ... .__getchild__(blah_n)

  • None of the nodes between bar and foo have the same label label_str.

Note: A node does not count as its own ancestor.

Specification and implementation of ancestry tracking

NOTE: This section is outdated and needlessly complicated

tl;dr: Visible ancestry is communicated from parent to child during .__getchild__(...). Shallow copies of the child are made, as necessary, to ensure that the same node cannot be in two different pytrees (or in two different locations of the same pytree) with conflicting visible ancestries.

  • The visible ancestors of treepie nodes are tracked in a central database. It can be thought of as a two-level nested dict: _internal_ancestors_database[descendant][label_str] returns a visible ancestor of descendant with label label_str (assuming there is one).
    Sidenote: The fact that descendent may not be hashable is handled appropriately internally. Lookup for the first level of the dictionary is by the actual identity of descendant (think of it as an is check), not by equality. So, copies of a node don’t share the same entry in the database.

  • When a treepie node is first created and initialized, it has no ancestors listed in the database. This also applies to nodes created during pytree unflattening (which can happen when nodes are sent through JAX-transforms, when nodes are acted on by jax.tree.map, etc.)

  • When a treepie node is assigned as a child-candidate of a parent treepie node, it is stored within the parent (indexed by the corresponding attr-name, map-key, or sequence-index).

  • When a child is accessed via parent.__getchild__(...):

    1. The corresponding child-candidate is retrieved. It isn’t returned yet. The return to this access happens in step 3.

    2. Next, we compute what the ancestry of a child of parent should be. This is done by first retrieving the visible labeled ancestors of parent from the database, and adding parent itself to this group if it is labeled (overriding other other ancestors if labels clash). This can be thought of as

      what_ancestry_of_child_should_be = (
    _internal_ancestors_database[parent]
    | {label_str: parent for label_str in (parent.__node_labels__ or {})}
)

    3. Next, we check if _internal_ancestors_database[child] matches what_ancestry_of_child_should_be, element-wise, i.e., we check if they contain (a) the same label-strings and (b) identical values corresponding to each label-string.

      • If yes, the child-candidate is returned for the access-query.

      • If no, a shallow copy of child-candidate, say child_copy, is created. An entry is added in the ancestors database for child_copy. This can be thought of as

        _internal_ancestors_database[child_copy] = what_ancestry_of_child_should_be

        child_copy is set as the new candidate-child, for subsequent accesses. Finally, child_copy is returned for the access-query.

< ONLY WRITTEN UPTO HERE >

Finer implementation details

f

What does this mean in practice

What you set may not be what you get:

Setup:

  1. Let 'attr_name' be declared as a Node on a treepie dataclass, and let parent be an instance of that dataclass. Let child_candidate be a treepie node (i.e., a dataclass instance, or an instance of treepie.FrozenMap or treepie.FrozenSeq).

  2. Consider an assignment of the form parent.attr_name = child_candidate or equivalently setattr(parent, 'attr_name', child_candidate).

  3. After the assignment, let the child node be accessed as parent.__getchild__("attr_name") or as parent.attr_name (which implicitly calls parent.__getchild__("attr_name") before unwrapping the retrieved child, if it is composite). This

This access might return child_candidate. Or it might return a shallow copy of child candidate. That may have been made

When is foo a known ancestor of bar with label label_str?

  • Condition 1: foo is a treepie dataclass instance with label_str as one of its labels.

  • Condition 2: bar is a treepie node (i.e., a treepie dataclass instance, or an instance of treepie.FrozenMap or treepie.FrozenSeq).

  • Condition 3: The reference to bar was created by starting from foo and explicitly walking down a pytree, through an uninterrupted sequence of treepie nodes. Each step in this walk is an access of a child node from it’s parent as parent.__getchild__(<child_name_or_key_or_idx>). Note that this is equivalent to tpy.get_children_map(parent, <child_name_or_key_or_idx>)

  • tree a sequence of node treepie nodes, .

  • Unless the reference to the object bar was retrieved by traversing from foo to bar through a sequence of treepie nodes, each accessed from it’s parent

  • assert foo == tpy.get_labeled_ancestor(bar, label_str) will be succeed if and only if assert

Let 'attr_name' be declared as a Node on a treepie dataclass, and let parent be an instance of that dataclass. Let child_candidate be a treepie node (i.e., dataclass instance, or an instance of treepie.FrozenMap or treepie.FrozenSeq).

What happens during assignment: Consider an assignment of the form parent.attr_name = child_candidate or equivalently setattr(parent, 'attr_name', child_candidate). Since treepie dataclasses have emulated immutability, such an assignment will likely be within an __init__ method as self.attr_name = child_candidate. During such an assignment, child_candidate gets stored within parent as a “possibly unsynched child” corresponding to the name attr_name. That’s it.

What happens during access: When a child is accessed as parent.__getchild__('attr_name'), the parent’s “possibly unsynced child” corresponding to that name, i.e, child_candidate, is retrieved. We then go through the following flowchart:

  1. Question 1: Are parent and child_candidate in sync with each other? If yes, return child_candidate. If no, proceed to the next step.

  2. Question 2: Do they need to sync up? If no, return child_candidate. If yes, proceed to the next step.

  3. Make a shallow copy of child_candidate, say child_copy. Store child_copy as the new “possibly unsynced child” and proceed to the next step.

  4. Question 3: Do parent and child_copy need to sync up? If no, return child_copy. If yes, proceed to the next step.

  5. Sync parent with child_copy and return child_copy.

What does it mean to sync two nodes? Any two nodes are unsynced by default. A parent and a child node can become synced only when step 5 happens. Syncing applies to specific objects. foo and bar being synced doesn’t mean that copy.copy(foo) and bar are synced. During syncing, a central database storing the ancestors of nodes gets updated. The known labeled ancestors of parent (along with parent itself, if it is labeled) will be associated with child_copy in this database. The database query is by identity of the descendant, so ancestors of foo are not automatically the ancestors of copy.copy(foo). Also, to be explicit, syncing is directional, i.e, parent–child relationship is not symmetric.

When do two unsynced nodes not need to sync up? If parent is unlabeled and has no known labeled ancestors AND child_candidate hasn’t synced up with any node, then parent and child_candidate don’t need to sync up, because there’s correct ancestry to communicate or wrong ancestry to erase.

Note that in step 3, we only make a shallow copy. Also note that even after syncing up in step 5, child_copy is only stored as a “possibly unsynced child” within parent. The check in step 1 will be performed during subsequent accesses.

What does it mean to be a “known labeled ancestor”: If there is a

That’s a wall of text. Here’s the important tl;dr

  • After parent.attr_name = child, child doesn’t necessarily become a child node of parent. A shallow copy of it might be the child.

  • How do you know if foo is an ancestor of bar? Can bar be reached from foo as follows:

    assert bar == foo.__getchild__(<blah_1>).__getchild__(<blah_2>)....__getchild__(<blah_n>)

    where each intermediate node is a treepie node? If yes, foo is an ancestor of bar. If no, then it isn’t.

Note

Here’s the cognitive load you take on when using the ancestry awareness feature.

After parent.attr_name = child:

  • tpy.get_labeled_ancestor(child, label) will not return the relevant ancestors. What you want to do instead is tpy.get_labeled_ancestor(parent.__getchild__('attr_name'), label).

  • In typical use cases, you would just be doing tpy.get_labeled_ancestor(self, label) from some method (or tpy.get_labeled_ancestor(parent, label) from an __unwrap__ method). The previous point is irrelevant for such cases.

  • assert parent.__getchild__('attr_name') is child may or may not fail (because of ancestry awareness). Try not to rely on such is checks.

  • assert parent.attr_name is child may or may not fail (because of ancestry awareness and/or composite nodes). Try not to rely on such is checks.

  • Mutating child may or may not affect the child-candidate stored inside parent. But, you shouldn’t be mutating a logically immutable object in the first place.

  • Mutating parent.get_child("attr_name") may or may not affect child. But, you shouldn’t be mutating a logically immutable object in the first place.

Hopefully these aren’t too bad.

This is the all the cognitive load

Let’s understand the ancestry semantics by probing these questions. For simplicity, let’s just use the following three classes.

class Descendant(tpy.Dataclass): ...


class LabeledAncestor(tpy.Dataclass):
    __node_labels__ = ("label",)
    desc: Node[Descendant]

    def __init__(self, desc: Descendant) -> None:
        self.desc = desc


class UnlabeledAncestor(tpy.Dataclass):
    desc: Node[Descendant]

    def __init__(self, desc: Descendant) -> None:
        self.desc = desc

We’ve defined the __init__ method on the ancestors just to show that nothing special needs to be done if overriding the __init__ method autogenerated by the dataclass-mechanism.

Now let’s create a candidate descendant object and place it within a pytree.

desc_candidate = Descendant()

labeled_ancestor = LabeledAncestor(desc_candidate)

Who is desc_candidate’s ancestor? Answer: Nobody.

tpy.get_labeled_ancestor(desc_candidate, "label")

RuntimeError: 'node' is not *aware* of any labeled ancestors.

However, labeled_ancestor.__getchild__("desc") does have labeled_ancestor as its ancestor:

assert labeled_ancestor == tpy.get_labeled_ancestor(
    labeled_ancestor.__getchild__("desc"), "label"
)

In the current example, only because desc_candidate is not a composite node (i.e., doesn’t have an __unwrap__ method), labeled_ancestor.__getchild__("desc") is equivalent to labeled.desc. So the following also works:

assert labeled_ancestor == tpy.get_labeled_ancestor(labeled_ancestor.desc, "label")

Something happens under the hood, so that labeled_ancestor.__getchild__("desc") is not desc_candidate. The specification of what happens under the hood is as follows:

Consider an assignment of the following form: parent.attr_name = child_candidate, where parent is an instance of a treepie dataclass. Since treepie nodes emulate immutability, this assignment will likely be happening inside an __init__ method as self.attr_name = child_candidate.

Defn 1: A parent that has ancestry to broadcast

Defn 2: A node that has been synced with it’s parent

If the following conditions are ALL satisfied, then parent.__getchild__("attr_name") will be a shallow copy of child_candidate instead of child_candidate itself:

  1. child_candidate is an instance of treepie.Dataclass, treepie.FrozenSeq, treepie.FrozenMap, treepie.FrozenSeqLeaf, or treepie.FrozenMapLeaf.

  2. attr_name has been declared as a Node of parent (as opposed to Aux, InitOnly, etc.)

  3. parent and child_candidate aren’t in sync with each other already.

So, after parent.attr_name = child_candidate, running assert parent.__getchild__("attr_name") is child_candidate might fail sometimes. This is the magic/caveat users need to keep in mind.

Note: assert parent.attr_name is child_candidate might fail even if we don’t have ancestry broadcasting, if child_candidate is a composite node.

What does being in sync (in condition 3 above) mean? Let parent be an instance of a treepie dataclass. Let child_candidate be an instance of a treepie node.

  1. If child_candidate is was already a child node of parent, then they are in sync. In other words:

    class Bar(tpy.Dataclass): ...

class Foo(tpy.Dataclass):
    __node_labels__ = ("label",)
    bar: Node[Bar]

def __init__

However, assert parent.__getchild__("attr_name") is parent.__getchild__("attr_name") will always succeed.

assert labeled_ancestor_1 is not labeled_ancestor_2
assert tpy.get_labeled_ancestor(labeled_ancestor_1.desc, "label") is labeled_ancestor_1
assert tpy.get_labeled_ancestor(labeled_ancestor_2.desc, "label") is labeled_ancestor_2

NameError: name 'labeled_ancestor_1' is not defined