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
- Field Declaration: default values for class variables that are set using the
- 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
node=Falseit 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()