Skip to content

PyTree utilities¤

The jaxnasium.tree package provides convenience pytree functions used in the various RL algorithms which aren't found in the used higher-level libraries (equinox / jax / optax).

import jaxnasium as jym

jym.tree.mean(pytree)

Reductions¤

jaxnasium.tree.mean(tree) ¤

Computes the global mean of the leaves of a pytree.

jaxnasium.tree.sum(tree: Any, axis: int | tuple[int, ...] | None = None) ¤

Compute the sum of all the elements in a pytree If axis is provided, sums each leaf over the specified axis and then adds adds the resulting leafs.

jaxnasium.tree.batch_mean(values, batch_axes: int | tuple[int, ...] = 0) ¤

Average over all non-batch axes of each leaf in a pytree, and across leaves. The batch axes(s) is/are assumed to be the leading axes.

All elements are weighted equally; essentially this is then a jaxnasium.tree.sum / .

Identical to tree_batch_sum up to a division by the (static) number of reduced elements.

Arguments: values: Pytree of JAX arrays. Every leaf must have at least len(batch_axes) leading dimensions. batch_axes: Leading axes to exclude from the reduction.

Returns: A JAX array with the same shape as the batch dimensions.

Notes: - For a single leaf with only batch dimensions, this is a no-op.

Example: >>> tree = { ... "a": jnp.array([[1.0, 2.0], [3.0, 4.0]]), ... "b": jnp.array([[5.0, 6.0], [7.0, 8.0]]), ... } >>> tree_batch_mean(tree, batch_axes=0) Array([3.5, 5.5])

jaxnasium.tree.batch_sum(values, batch_axes: int | tuple[int, ...] = 0) ¤

Sum over all non-batch axes of each leaf in a pytree, then sum (reduce) across leaves. The batch axes(s) is/are assumed to be the leading axes.

This is essentially jaxnasium.tree.sum or optax.tree.sum but with a variable axis argument resulting in a sum over all non-batch axes.

Arguments: values: Pytree of JAX arrays. Every leaf must have at least len(batch_axes) leading dimensions. batch_axes: Leading axes to exclude from the sum.

Returns: A JAX array with the same shape as the batch dimensions.

Notes: - For a single leaf with only batch dimensions, this is a no-op.

Example: >>> tree = {"a": jnp.array([[1, 2], [3, 4]]), "b": jnp.array([[5, 6], [7, 8]])} >>> tree_batch_sum(tree, batch_axes=0) Array([14, 22])

>>> tree2 = {"x": jnp.ones((2, 3, 4)), "y": jnp.ones((2, 3, 4))}
>>> tree_batch_sum(tree2, batch_axes=(0, 1))
Array([[8., 8., 8.], [8., 8., 8.]])

Arithmetic¤

jaxnasium.tree.add(tree_A: PyTree, tree_B_or_prefix: PyTree | float | Array) -> PyTree ¤

Add two pytrees or add a scalar, array, or prefix-pytree to each leaf of a pytree.

Arguments: tree_A: First pytree. tree_B_or_prefix: Second pytree or scalar, array, or prefix-pytree of tree_A.

Example: 

    >>> tree_A = [5, 6]
    >>> tree_B = [10, 11]
    >>> tree_add(tree_A, tree_B)
    [Array(15, dtype=int32), Array(17, dtype=int32)]

    >>> tree_A = {'a': jnp.array([1, 2]), 'b': jnp.array(3)}
    >>> tree_B = 1
    >>> tree_add(tree_A, tree_B)
    {'a': Array([2, 3], dtype=int32), 'b': Array(4, dtype=int32)}

jaxnasium.tree.mul(tree_A: PyTree, tree_B_or_prefix: PyTree | float | Array) -> PyTree ¤

Multiply two pytrees or multiply a scalar, array, or prefix-pytree to each leaf of a pytree.

Arguments: tree_A: First pytree. tree_B_or_prefix: Second pytree or scalar, array, or prefix-pytree of tree_A.

Example: 

>>> tree_A = [5, 6]
>>> tree_B = [10, 11]
>>> tree_mul(tree_A, tree_B)
[Array(50, dtype=int32), Array(66, dtype=int32)]

>>> tree_A = {'a': jnp.array([1, 2]), 'b': jnp.array(3)}
>>> tree_B = 2
>>> tree_mul(tree_A, tree_B)
{'a': Array([2, 4], dtype=int32), 'b': Array(6, dtype=int32)}

jaxnasium.tree.clip(tree: PyTree, min_value: float | Array, max_value: float | Array) -> PyTree ¤

Clip the leaves of a pytree to a specified range.

Arguments: tree: A pytree whose leaves are array-like. min_value: Minimum value to clip to (scalar or array). max_value: Maximum value to clip to (scalar or array).

Structure¤

jaxnasium.tree.get_first(tree: PyTree, key: str) -> Any ¤

Get the first value from a pytree with the given key.

Arguments:

  • tree: A pytree.
  • key: A string key.

Returns: The first value from the pytree with the given key.

Raises: KeyError: If the key is not found in the pytree.

jaxnasium.tree.map_one_level(fn: Callable, tree, *rest) ¤

Simple jax.tree.map operation over the first level of a pytree.

Arguments:

  • fn: A function to map over the pytree.
  • tree: A pytree.
  • *rest: Additional pytrees to map over.

jaxnasium.tree.map_distribution(fn: Callable, tree, *rest) ¤

Map a function with distrax.Distribution instances marked as leaves. Additionally, if one of the inputs is a DistraxContainer, the function is applied to the distribution attribute of the DistraxContainer.

Arguments:

  • fn: A function to map over the pytree.
  • tree: A pytree.
  • *rest: Additional pytrees to map over.

jaxnasium.tree.stack(pytrees: PyTree, *, axis=0) -> PyTree ¤

Stack corresponding leaves of pytrees along the specified axis.

Interprets the root node's immediate children as a batch of N pytrees that all share the same structure. For each leaf, stacks the N leaves along axis using jnp.stack. This does not traverse deeper than one level when determining what to stack.

Arguments:

  • pytrees: A pytree whose root has N immediate children. Each child must have the same pytree structure. Corresponding leaves must be array-like and have identical shapes and dtypes (compatible with jnp.stack).
  • axis: Axis along which to insert the new dimension of size N in each stacked leaf (default=0).

Returns: A pytree with the same structure as a single direct-child element of pytrees, where each leaf is the stack of the corresponding leaves across all elements, with a new dimension of size N inserted at axis.

Example: 

    >>> trees = (
    ...     [jnp.array([1, 2]), jnp.array(4)],
    ...     [jnp.array([5, 5]), jnp.array(3)],
    ... )
    >>> stack_one_level(trees, axis=0)
    [Array([[1, 2], [5, 5]], dtype=int32), Array([4, 3], dtype=int32)]
Evolved from: link.

jaxnasium.tree.unstack(tree, *, axis=0, structure: PyTreeDef | None = None) ¤

Inverse of stack: split a pytree whose leaves were stacked along axis into N separate pytrees.

If structure is provided (e.g., from eqx.tree_flatten_one_level), the list of N pytrees is immediately placed back into that container and returned as a single pytree.

Arguments:

  • tree: A pytree whose leaves are array-like and all share the same size N along axis.
  • axis: The axis that carries the size-N dimension in each leaf (default=0).
  • structure: Optional PyTreeDef. If provided, the list of N pytrees is immediately placed back into that container and returned as a single pytree.

Returns: If structure is None: a list of N pytrees. Otherwise: a single pytree produced by unflattening structure with those N pytrees.

Example: 

    >>> trees = (
    ...     [jnp.array([1, 2]), jnp.array(4)],
    ...     [jnp.array([5, 5]), jnp.array(3)],
    ... )
    >>> batched = stack(trees, axis=0)
    >>> unstack(batched, axis=0)
    [[Array([1, 2], dtype=int32), Array(4, dtype=int32)],
     [Array([5, 5], dtype=int32), Array(3, dtype=int32)]]

jaxnasium.tree.concatenate(pytrees: PyTree, *, axis=0) -> PyTree ¤

Concatenate corresponding leaves of pytrees along the specified axis.

Interprets the root node's immediate children as a batch of N pytrees that all share the same structure. For each leaf, concatenates the N leaves along axis using jnp.concatenate. This does not traverse deeper than one level when determining what to concatenate.

Arguments:

  • pytrees: A pytree whose root has N immediate children. Each child must have the same pytree structure. Corresponding leaves must be array-like and compatible with jnp.concatenate along axis.
  • axis: Axis along which to concatenate corresponding leaves (default=0).

Returns: A pytree with the same structure as a single direct-child element of pytrees, where each leaf is the concatenation of the corresponding leaves across all elements.

Example: 

    >>> trees = (
    ...     [jnp.array([1, 2]), jnp.array([4])],
    ...     [jnp.array([5, 5]), jnp.array([3])],
    ... )
    >>> tree_concatenate(trees, axis=0)
    [Array([1, 2, 5, 5], dtype=int32), Array([4, 3], dtype=int32)]

jaxnasium.tree.ravel(tree: PyTree) -> Array ¤

Flatten a pytree of arrays into a single 1D array.

Arguments:

  • tree: A pytree whose leaves are array-like.

Returns: A 1D array containing every element of the pytree.

Example: 

    >>> tree = {'a': jnp.array([1, 2]), 'b': jnp.array(3)}
    >>> tree_ravel(tree)
    Array([1, 2, 3], dtype=int32)

jaxnasium.tree.gather_actions(tree: PyTree, actions: PyTree, axis=-1) ¤

Given a (pytree of) array-like values, gather the elements based on the indices provided in actions. If the arrays in tree are of the same shape as actions, the tree is assumed to be array of actions taken and the array is returned as is. This may be the case in continuous action spaces.

For example, when given a (pytree of) q-values for each possible action, this function will return the q-values corresponding to the actions taken. In continuous action spaces, q-values cannot be generated per action and tree will already contain the q-value for the action taken. This q-value is then returned as is. This is also infered when the indices are floating point indices.

Arguments: tree: Array or Pytree of arrays. actions: Array or same-structure Pytree of arrays as tree. The final axis of actions must contain elements that are valid indices for the corresponding arrays in tree. axis: The axis on which the actions live on each leaf within the tree.

Creation¤

jaxnasium.tree.zeros_like(tree: PyTree, dtype: DTypeLike | None = None) -> PyTree ¤

Creates an all-zeros PyTree with the same structure as tree.

Arguments: tree: A pytree. dtype: The dtype of the tree of zeros.

jaxnasium.tree.ones_like(tree: PyTree, dtype: DTypeLike | None = None) -> PyTree ¤

Creates an all-ones PyTree with the same structure as tree.

Arguments: tree: A pytree. dtype: The dtype of the tree of ones.

jaxnasium.tree.split_key_like(key: PRNGKeyArray, tree: PyTree, is_leaf: None | Callable[[Any], bool] = None) ¤

Split a JAX PRNGKey into a pytree of keys with the same structure as tree.

Arguments: key: A PRNGKeyArray to be split. tree: A pytree whose structure will be used to determine the number of keys to split. is_leaf: A function that determines which nodes in the tree are leaves.

jaxnasium.tree.split_key_like_structure(key: PRNGKeyArray, structure: PyTreeDef) ¤

Split a JAX PRNGKey into a pytree of keys with the same structure as structure.

Similar to optax.tree_utils.tree_split_key_like, but operates on PyTreeDefs.

Arguments: key: A PRNGKeyArray to be split. agent_structure: A pytree structure of agents.