Mismatched Zip API Reference

Zip and map a function over multiple PyTrees that have mismatched structure or shape.

Comparison

• JAX: Extends jax.tree_util.tree_map (when mapping over multiple trees), but handles mismatched structures (different dictionary keys, different list/tuple lengths) via intersection or padding rather than raising shape/structure errors. Supports broadcasting of scalar leaves to the whole tree structure.

Example

import jaxtree_extensions as jte

tree1 = {"a": [1, 2], "b": 3}
tree2 = {"a": [10, 20, 30], "c": 4}

# Zip matching elements via intersection
res = jte.tree_zip_mismatched(
    lambda x, y: x + y,
    tree1,
    tree2,
    mode="intersection"
)
# res = {"a": [11, 22]}

# Zip with padding and fallback fill value
res_pad = jte.tree_zip_mismatched(
    lambda x, y: x + y,
    tree1,
    tree2,
    mode="padding",
    fill_value=0
)
# res_pad = {"a": [11, 22, 30], "b": 3, "c": 4}

jaxtree_extensions.tree_zip_mismatched(f: collections.abc.Callable[..., typing.Any], tree: typing.Any, *rest: typing.Any, mode: str = 'intersection', fill_value: typing.Any = None, is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> typing.Any

Zip multiple PyTrees by mapping f over their leaves, handling mismatches.

If structures mismatch, the behavior is determined by the mode parameter. If a leaf is encountered in one tree and a container in another, the leaf is broadcasted to match the container's structure.

Parameters:

Name	Type	Description	Default
f	Callable[..., Any]	The function to map over aligned leaves.	required
tree	Any	The primary PyTree.	required
*rest	Any	Additional PyTrees to align and map.	required
mode	str	How to handle mismatched keys or elements: - "intersection": Drop unmatched keys or elements from lists/tuples. - "padding": Pad unmatched elements/branches using fill_value.	'intersection'
fill_value	Any	The value used to fill missing elements in "padding" mode.	None
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	The zipped PyTree.