Jaxtree Extensions
jaxtree-extensions extends JAX's native jax.tree_util with extra tree utilities for JAX, Equinox, and other pytree frameworks.
Core Utilities
| Function | Description | Comparison |
|---|---|---|
tree_map_prune |
Map and filter leaves by predicate, collapsing containers. | Similar to jax.tree_util.tree_map but drops/collapses leaves. |
tree_select_by_path / tree_at_path |
Mask or edit nodes matching wildcard path strings (e.g. **/bias). |
Similar to equinox.tree_at but using glob/wildcard patterns. |
tree_zip_mismatched |
Zip mismatched trees via intersection/padding and broadcast scalars. | Extends jax.tree_util.tree_map to handle structure/shape mismatches. |
tree_map_enumerate |
Map over a tree passing a sequential flat index to each leaf. | Similar to jax.tree_util.tree_map but provides automatic indices 0, 1, 2... |
tree_compare |
Compare two trees and generate a detailed mismatch diagnostic report. | Similar to chex.assert_trees_all_close but returns a diagnostic TreeDiff. |
tree_reduce_with_path |
Reduce a tree passing key-path tuples of each leaf. | Extends jax.tree_util.tree_reduce with path-aware context. |
Installation
Install the package via PyPI:
pip install jaxtree-extensions
Or install locally for development:
uv pip install -e .
Quick Start
import jaxtree_extensions as jte
import jax.numpy as jnp
# Pruning elements
pruned, _ = jte.tree_map_prune(lambda x: x, {"a": [1, 2], "b": 3}, keep_fn=lambda x: x > 1)
# pruned = {"a": [2], "b": 3}
# Path-based editing
model = {"layer1": {"bias": jnp.ones(2)}, "layer2": {"bias": jnp.ones(2)}}
zero_bias = jte.tree_at_path(model, "**/bias", replace=0.0)
Running Tests
To run the test suite:
uv run pytest
1 maintainer
Metadata
1 maintainer
Release files for jaxtree-extensions 0.2.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | Actions |
|---|---|---|---|---|
| jaxtree_extensions-0.2.0.tar.gz | 117.7 kB | Without provenance | Details for jaxtree_extensions-0.2.0.tar.gz |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Actions Reset |
|---|---|---|---|---|
| jaxtree_extensions-0.2.0-py3-none-any.whl | Python 3 | none | any | Details for jaxtree_extensions-0.2.0-py3-none-any.whl |
Total release size: 131.5 kB
Release files / jaxtree_extensions-0.2.0.tar.gz
| Download URL | jaxtree_extensions-0.2.0.tar.gz |
|---|---|
| Size | 117.7 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
a3281851594e635846f0d6ccbc859c18b944c43a142d72204cae94e50b6756d3
|
|
BLAKE2b-256 checksum How to use checksums |
08208e3ec3479ef3e4a90ccf8a678801fd2e80fa931835ae633353668b34fd2c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.11.14 {"installer":{"name":"uv","version":"0.11.14","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Release files / jaxtree_extensions-0.2.0-py3-none-any.whl
| Download URL | jaxtree_extensions-0.2.0-py3-none-any.whl |
|---|---|
| Size | 13.8 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
8e5406e48f442135cbe1346aa742368a195fab13fb4f067010a81067b2d8a682
|
|
BLAKE2b-256 checksum How to use checksums |
a88d3a484e500c1f0618cbf8d3ef999232fde2d8a248715e7c38f65e1ea949f9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
uv/0.11.14 {"installer":{"name":"uv","version":"0.11.14","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
|
Security for jaxtree-extensions 0.2.0
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
