<commit_before>Pytrees
========

What is a pytree?
^^^^^^^^^^^^^^^^^

In JAX, we use the term *pytree* to refer to a tree-like structure built out of
container-like Python objects. Classes are considered container-like if they
are in the pytree registry, which by default includes lists, tuples, and dicts.
That is:

1. any object whose type is *not* in the pytree container registry is
   considered a *leaf* pytree;
2. any object whose type is in the pytree container registry, and which
   contains pytrees, is considered a pytree.

For each entry in the pytree container registry, a container-like type is
registered with a pair of functions which specify how to convert an instance of
the container type to a ``(children, metadata)`` pair and how to convert such a
pair back to an instance of the container type. Using these functions, JAX can
canonicalize any tree of reigstered container objects into tuples.

Example pytrees::

  [1, "a", object()]  # 3 leaves

  (1, (2, 3), ())  # 3 leaves

  [1, {"k1": 2, "k2": (3, 4)}, 5]  # 5 leaves

JAX can be extended to consider other container types as pytrees; see
`Extending pytrees`_ below.

Pytrees and JAX functions
^^^^^^^^^^^^^^^^^^^^^^^^^

Many JAX functions, like ``jax.lax.scan``, operate over pytrees of arrays.
JAX function transformations can be applied to functions that accept as input
and produce as output pytrees of arrays.

Applying optional parameters to pytrees
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Some JAX function transformations take optional parameters that specify how
certain input or output values should be treated (e.g. the ``in_axes`` and
``out_axes`` arguments to ``vmap``). These parameters can also be pytrees, and
their structure must correspond to the pytree structure of the corresponding
arguments. In particular, to be able to "match up" leaves in these parameter
pytrees with values in the argument pytrees, the parameter pytrees are often
constrained to be tree prefixes of the argument pytrees.

For example, if we pass the following input to ``vmap`` (note that the input
arguments to a function considered a tuple)::

  (a1, {"k1": a2, "k2": a3})

We can use the following ``in_axes`` pytree to specify that only the ``k2``
argument is mapped (``axis=0``) and the rest aren't mapped over
(``axis=None``)::

  (None, {"k1": None, "k2": 0})

The optional parameter pytree structure must match that of the main input
pytree. However, the optional parameters can optionally be specified as a
"prefix" pytree, meaning that a single leaf value can be applied to an entire
sub-pytree. For example, if we have the same ``vmap`` input as above, but wish
to only map over the dictionary argument, we can use::

  (None, 0)  # equivalent to (None, {"k1": 0, "k2": 0})

Or, if we want every argument to be mapped, we can simply write a single leaf
value that is applied over the entire argument tuple pytree::

  0

This happens to be the default ``in_axes`` value for ``vmap``!

The same logic applies to other optional parameters that refer to specific input
or output values of a transformed function, e.g. ``vmap``'s ``out_axes``.


Developer information
^^^^^^^^^^^^^^^^^^^^^^

*This is primarily JAX internal documentation, end-users are not supposed to need
to understand this to use JAX, except when registering new user-defined<commit_msg>Add description how to view pytree definition<commit_after>Pytrees
========

What is a pytree?
^^^^^^^^^^^^^^^^^

In JAX, we use the term *pytree* to refer to a tree-like structure built out of
container-like Python objects. Classes are considered container-like if they
are in the pytree registry, which by default includes lists, tuples, and dicts.
That is:

1. any object whose type is *not* in the pytree container registry is
   considered a *leaf* pytree;
2. any object whose type is in the pytree container registry, and which
   contains pytrees, is considered a pytree.

For each entry in the pytree container registry, a container-like type is
registered with a pair of functions which specify how to convert an instance of
the container type to a ``(children, metadata)`` pair and how to convert such a
pair back to an instance of the container type. Using these functions, JAX can
canonicalize any tree of reigstered container objects into tuples.

Example pytrees::

  [1, "a", object()]  # 3 leaves

  (1, (2, 3), ())  # 3 leaves

  [1, {"k1": 2, "k2": (3, 4)}, 5]  # 5 leaves

JAX can be extended to consider other container types as pytrees; see
`Extending pytrees`_ below.

Pytrees and JAX functions
^^^^^^^^^^^^^^^^^^^^^^^^^

Many JAX functions, like ``jax.lax.scan``, operate over pytrees of arrays.
JAX function transformations can be applied to functions that accept as input
and produce as output pytrees of arrays.

Applying optional parameters to pytrees
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Some JAX function transformations take optional parameters that specify how
certain input or output values should be treated (e.g. the ``in_axes`` and
``out_axes`` arguments to ``vmap``). These parameters can also be pytrees, and
their structure must correspond to the pytree structure of the corresponding
arguments. In particular, to be able to "match up" leaves in these parameter
pytrees with values in the argument pytrees, the parameter pytrees are often
constrained to be tree prefixes of the argument pytrees.

For example, if we pass the following input to ``vmap`` (note that the input
arguments to a function considered a tuple)::

  (a1, {"k1": a2, "k2": a3})

We can use the following ``in_axes`` pytree to specify that only the ``k2``
argument is mapped (``axis=0``) and the rest aren't mapped over
(``axis=None``)::

  (None, {"k1": None, "k2": 0})

The optional parameter pytree structure must match that of the main input
pytree. However, the optional parameters can optionally be specified as a
"prefix" pytree, meaning that a single leaf value can be applied to an entire
sub-pytree. For example, if we have the same ``vmap`` input as above, but wish
to only map over the dictionary argument, we can use::

  (None, 0)  # equivalent to (None, {"k1": 0, "k2": 0})

Or, if we want every argument to be mapped, we can simply write a single leaf
value that is applied over the entire argument tuple pytree::

  0

This happens to be the default ``in_axes`` value for ``vmap``!

The same logic applies to other optional parameters that refer to specific input
or output values of a transformed function, e.g. ``vmap``'s ``out_axes``.

Viewing the pytree definition of an object
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
To view the pytree definition of an arbitrary ``object`` for debugging purposes, you can use::

  from jax.tree_util import tree_structure
  print(tree_structure(object))

Developer information
^^^^^^^^^^^^^^^^^^^^^^

*This is primarily JAX internal documentation, end-users are not supposed to need
to understand this to use JAX, except when registering new user-defined
{'fraction_non_alphanumeric': 0.10713772799080856, 'fraction_numerical': 0.009191440471061324, 'mean_word_length': 4.666395443449959, 'pattern_counts': {'":': 16, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3049107', 'n_tokens_mistral': 2005, 'n_tokens_neox': 1935, 'n_words': 1044}