Skip to content

User Guide

Module is the main construct in Treex, it inherits from treeo.Tree and adds a couple of convenient methods. We recommend that you review the core concepts of Treeo but we will provide a brief overview.


These are the core concepts from Treeo:

  • Type Annotation: (type hints) types you set while defining a variable after the : symbol.
  • Field Declaration: default values for class variables that are set using the field function.
  • Node Field: A field that is declared as a node, that is, its content is part of the tree leaves.
  • Static Field: A field that is declared as a static, that is, its content is not part of the leaves.
  • Field Kind: An associated type, separate from the type annotation, that gives semantic meaning to the field.

In code these terms map to the following:

class MyModule(tx.Module):
    #  field      annotation   ------------declaration---------------
    #    v            v        v                                    v
    some_field : jnp.ndarray = to.field(node=True, kind=tx.Parameter)
    #                                      ^                ^
    #                                 node status       field kind
Here if node=False it would mean that the field is static, else is a node. The previous is written more compactly as:

class MyModule(tx.Module):
    some_field: jnp.ndarray = tx.Parameter.node()