Path Selection API Reference¤

Select and modify nodes in a PyTree using glob/wildcard path pattern matching.

Comparison¤

Equinox: Similar to equinox.tree_at, which replaces values/calls functions at specified nodes, but supports wildcard selectors (*, **, ?) instead of static lambda accessors, allowing bulk modification of attributes across layers.
JAX: Complements JAX's key paths, offering wildcard matching over path string structures.

Example¤

import jaxtree_extensions as jte
import jax.numpy as jnp

model = {
    "layer1": {"weight": jnp.ones((2, 2)), "bias": jnp.ones(2)},
    "layer2": {"weight": jnp.ones((2, 2)), "bias": jnp.ones(2)},
}

# Select all biases in all layers
mask = jte.tree_select_by_path(model, "**/bias")
# mask = {
#     "layer1": {"weight": False, "bias": True},
#     "layer2": {"weight": False, "bias": True},
# }

# Zero out all biases
zero_bias_model = jte.tree_at_path(model, "**/bias", replace=0.0)

`jaxtree_extensions.tree_select_by_path(tree: typing.Any, pattern: str, is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> typing.Any` ¤

Return a boolean mask PyTree matching a glob-like path pattern.

Leaves matching the pattern are True, others are False.

Parameters:

Name	Type	Description	Default
`tree`	`Any`	The PyTree to inspect.	required
`pattern`	`str`	A glob-like path pattern (e.g. "**/weight").	required
`is_leaf`	`Callable[[Any], bool] \| None`	A predicate called on each node to determine if it should be treated as a leaf.	`None`

Returns:

Type	Description
`Any`	A boolean mask PyTree matching the original structure.

`jaxtree_extensions.tree_at_path(tree: typing.Any, pattern: str, replace: typing.Any = None, replace_fn: collections.abc.Callable[[typing.Any], typing.Any] | None = None, is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> typing.Any` ¤

Modify a PyTree by replacing leaves matching a path pattern.