Traversal

refrax.Traversal(tree, base_path, sub_paths)

Bases: Generic[TRoot]

Represents a multi-target focus within an immutable PyTree.

Applies mutations across all targets simultaneously using eqx.tree_at, making it fully compatible with JAX JIT compilation and tracing.

Parameters:

Name	Type	Description	Default
tree	TRoot	The root immutable object (e.g., Equinox module or dataclass).	required
base_path	list[PathStep]	The shared path from the root to the divergence point.	required
sub_paths	list[list[PathStep]]	A list of diverging paths, one for each targeted element.	required

Source code in refrax/traversal.py

def __init__(self, tree: TRoot, base_path: list[PathStep], sub_paths: list[list[PathStep]]) -> None:
    self._tree = tree
    self._base_path = base_path
    self._sub_paths = sub_paths

__getattr__(name)

Broadens the traversal by appending an attribute access to every currently focused target.

Parameters:

Name	Type	Description	Default
name	str	The name of the attribute to access on all targets.	required

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal focused one level deeper.

Source code in refrax/traversal.py

def __getattr__(self, name: str) -> "Traversal[TRoot]":
    """Broadens the traversal by appending an attribute access to every currently focused target.

    Args:
        name (str): The name of the attribute to access on all targets.

    Returns:
        Traversal[TRoot]: A new Traversal focused one level deeper.
    """
    new_sub_paths = [path + [("attr", name)] for path in self._sub_paths]
    return Traversal(self._tree, self._base_path, new_sub_paths)

__getitem__(key)

Broadens the traversal by appending an item/index access to every currently focused target.

Strings are treated as attributes.

Parameters:

Name	Type	Description	Default
key	Any	The index, dictionary key, or attribute string to access.	required

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal focused one level deeper.

Source code in refrax/traversal.py

def __getitem__(self, key: Any) -> "Traversal[TRoot]":
    """Broadens the traversal by appending an item/index access to every currently focused target. 

    Strings are treated as attributes.

    Args:
        key (Any): The index, dictionary key, or attribute string to access.

    Returns:
        Traversal[TRoot]: A new Traversal focused one level deeper.
    """
    op: PathOp = "attr" if isinstance(key, str) else "item"
    new_sub_paths = [path + [(op, key)] for path in self._sub_paths]
    return Traversal(self._tree, self._base_path, new_sub_paths)

apply(func)

Applies a function to all selected targets simultaneously.

Parameters:

Name	Type	Description	Default
func	Callable[[Any], Any]	The transformation function to apply to each target.	required

Returns:

Name	Type	Description
TRoot	TRoot	A new instance of the root tree with all targets updated.

Source code in refrax/traversal.py

def apply(self, func: Callable[[Any], Any]) -> TRoot:
    """Applies a function to all selected targets simultaneously.

    Args:
        func (Callable[[Any], Any]): The transformation function to apply to each target.

    Returns:
        TRoot: A new instance of the root tree with all targets updated.
    """
    if not self._sub_paths:
        return self._tree

    return eqx.tree_at(self._get_targets_from, self._tree, replace_fn=func)    

exclude(predicate)

Filters the currently focused targets, dropping those that match the condition.

This is the logical inverse of .filter().

Parameters:

Name	Type	Description	Default
predicate	Callable[[Any], bool]	A function that returns True to drop a target.	required

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal with the matching targets removed.

Source code in refrax/traversal.py

def exclude(self, predicate: Callable[[Any], bool]) -> "Traversal[TRoot]":
    """Filters the currently focused targets, dropping those that match the condition.

    This is the logical inverse of `.filter()`.

    Args:
        predicate (Callable[[Any], bool]): A function that returns True to drop a target.

    Returns:
        Traversal[TRoot]: A new Traversal with the matching targets removed.
    """
    return self.filter(lambda x: not predicate(x))

filter(predicate)

Filters the currently focused targets, keeping only those that match the condition.

Parameters:

Name	Type	Description	Default
predicate	Callable[[Any], bool]	A function that returns True to keep a target, or False to drop it.	required

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal focused only on the targets that passed the filter.

Source code in refrax/traversal.py

def filter(self, predicate: Callable[[Any], bool]) -> "Traversal[TRoot]":
    """Filters the currently focused targets, keeping only those that match the condition.

    Args:
        predicate (Callable[[Any], bool]): A function that returns True to keep a target, 
            or False to drop it.

    Returns:
        Traversal[TRoot]: A new Traversal focused only on the targets that passed the filter.
    """
    from refrax.lens import Lens

    filtered_paths = []
    for path in self._sub_paths:
        target_val = Lens(self._tree, self._base_path + path).get()
        if predicate(target_val):
            filtered_paths.append(path)

    return Traversal(self._tree, self._base_path, filtered_paths)

get()

Extracts all focused values.

Returns:

Type	Description
list[Any]	list[Any]: A list containing the values of all currently focused targets.

Source code in refrax/traversal.py

def get(self) -> list[Any]:
    """Extracts all focused values.

    Returns:
        list[Any]: A list containing the values of all currently focused targets.
    """
    return list(self._get_targets_from(self._tree))    

path(target_path)

Broadens the traversal by appending a string path or JAX KeyPath to every currently focused target.

Parameters:

Name	Type	Description	Default
target_path	str | tuple	The path string (e.g., '.res.R') or native JAX KeyPath tuple.	required

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal focused one level deeper across all branches.

Source code in refrax/traversal.py

def path(self, target_path: str | tuple) -> "Traversal[TRoot]":
    """
    Broadens the traversal by appending a string path or JAX KeyPath to every currently focused target.

    Args:
        target_path (str | tuple): The path string (e.g., '.res.R') or native 
            JAX KeyPath tuple.

    Returns:
        Traversal[TRoot]: A new Traversal focused one level deeper across all branches.
    """
    if isinstance(target_path, str):
        parsed_steps = parse_string_path(target_path)
    elif isinstance(target_path, tuple):
        parsed_steps = translate_jax_path(target_path)
    else:
        raise TypeError(f"Expected string or JAX path tuple, got {type(target_path).__name__}")

    # Append the parsed steps to all diverging paths
    new_sub_paths = [path + parsed_steps for path in self._sub_paths]
    return Traversal(self._tree, self._base_path, new_sub_paths)

prune(*paths)

Filters out targets from the Traversal that intersect with the given paths.

Crucially, this prunes a target if it is exactly the excluded path, inside the excluded path, OR a parent of the excluded path. (Mutating a parent implicitly mutates its children, so parents of excluded paths must be pruned).

Parameters:

Name	Type	Description	Default
*paths	str | tuple	String paths or JAX KeyPath tuples to protect from mutation.	()

Returns:

Type	Description
Traversal[TRoot]	Traversal[TRoot]: A new Traversal with the intersecting paths removed.

Source code in refrax/traversal.py

def prune(self, *paths: str | tuple) -> "Traversal[TRoot]":
    """
    Filters out targets from the Traversal that intersect with the given paths.

    Crucially, this prunes a target if it is exactly the excluded path, inside 
    the excluded path, OR a parent of the excluded path. (Mutating a parent 
    implicitly mutates its children, so parents of excluded paths must be pruned).

    Args:
        *paths (str | tuple): String paths or JAX KeyPath tuples to protect from mutation.

    Returns:
        Traversal[TRoot]: A new Traversal with the intersecting paths removed.
    """
    # Parse the user's paths into refrax PathSteps
    parsed_excludes = []
    for p in paths:
        if isinstance(p, str):
            parsed_excludes.append(parse_string_path(p))
        elif isinstance(p, tuple):
            parsed_excludes.append(translate_jax_path(p))
        else:
            raise TypeError(f"Expected string or JAX path tuple, got {type(p).__name__}")

    kept_sub_paths = []
    for path in self._sub_paths:
        # Reconstruct the absolute path from the root to the target
        full_target_path = self._base_path + path

        should_prune = False
        for ex_path in parsed_excludes:
            # 1. Is the target a PARENT of the excluded path? (or an exact match)
            if len(full_target_path) <= len(ex_path) and ex_path[:len(full_target_path)] == full_target_path:
                should_prune = True
                break

            # 2. Is the target a CHILD of the excluded path?
            if len(full_target_path) > len(ex_path) and full_target_path[:len(ex_path)] == ex_path:
                should_prune = True
                break

        if not should_prune:
            kept_sub_paths.append(path)

    return Traversal(self._tree, self._base_path, kept_sub_paths)

set(value)

Sets all selected targets to a specific value, or maps a sequence of values 1-to-1.

Parameters:

Name	Type	Description	Default
value	Any	The new value to assign. If value is a list/tuple of the exact same length as the focused targets, they will be mapped 1-to-1.	required

Returns:

Name	Type	Description
TRoot	TRoot	A new instance of the root tree with the updated values.

Source code in refrax/traversal.py

def set(self, value: Any) -> TRoot:
    """Sets all selected targets to a specific value, or maps a sequence of values 1-to-1.

    Args:
        value (Any): The new value to assign. If value is a list/tuple of the exact 
            same length as the focused targets, they will be mapped 1-to-1.

    Returns:
        TRoot: A new instance of the root tree with the updated values.
    """
    if not self._sub_paths:
        return self._tree

    # 1-to-1 Mapping: If the user passed a list of values matching our targets
    if isinstance(value, (list, tuple)) and len(value) == len(self._sub_paths):
        replacements = tuple(value)
    # Broadcasting: User passed a single value to apply everywhere
    else:
        replacements = tuple(value for _ in self._sub_paths)

    return eqx.tree_at(self._get_targets_from, self._tree, replace=replacements)