From ff800c79b0878b19c8fef956eca1a441aa7704d1 Mon Sep 17 00:00:00 2001 From: yspkm Date: Sat, 8 Aug 2026 21:17:31 +0800 Subject: [PATCH 1/4] Fix autoinit hashability metadata autoinit only synthesizes __init__, so eq_default must be false to preserve TreeClass's inherited hashability in static analysis. --- pytreeclass/_src/code_build.py | 2 +- tests/test_treeclass.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/pytreeclass/_src/code_build.py b/pytreeclass/_src/code_build.py index e64ca2f..7398f1a 100644 --- a/pytreeclass/_src/code_build.py +++ b/pytreeclass/_src/code_build.py @@ -482,7 +482,7 @@ def build_init_method(klass: type[T]) -> type[T]: return klass -@dataclass_transform(field_specifiers=(Field, field)) +@dataclass_transform(field_specifiers=(Field, field), eq_default=False) def autoinit(klass: type[T]) -> type[T]: """A class decorator that generates the ``__init__`` method from type hints. diff --git a/tests/test_treeclass.py b/tests/test_treeclass.py index 06cf29d..36a9296 100644 --- a/tests/test_treeclass.py +++ b/tests/test_treeclass.py @@ -623,6 +623,10 @@ def __init__(self, a): assert True +def test_autoinit_dataclass_transform_defaults(): + assert autoinit.__dataclass_transform__["eq_default"] is False + + def test_nohints(): assert convert_hints_to_fields(int) is int From c6005ef6306c40ee821a80af103a541f1e649592 Mon Sep 17 00:00:00 2001 From: yspkm Date: Sun, 9 Aug 2026 08:56:45 +0800 Subject: [PATCH 2/4] docs: use jax.tree_util.tree_map --- pytreeclass/_src/tree_mask.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/pytreeclass/_src/tree_mask.py b/pytreeclass/_src/tree_mask.py index bb68dd6..36526d9 100644 --- a/pytreeclass/_src/tree_mask.py +++ b/pytreeclass/_src/tree_mask.py @@ -205,9 +205,9 @@ def unfreeze(value: T) -> T: >>> frozen_value = tc.freeze(1) >>> tc.unfreeze(frozen_value) 1 - >>> # usage with `jax.tree_map` - >>> frozen_tree = jax.tree_map(tc.freeze, {"a": 1, "b": 2}) - >>> unfrozen_tree = jax.tree_map(tc.unfreeze, frozen_tree, is_leaf=tc.is_frozen) + >>> # usage with `jax.tree_util.tree_map` + >>> frozen_tree = jax.tree_util.tree_map(tc.freeze, {"a": 1, "b": 2}) + >>> unfrozen_tree = jax.tree_util.tree_map(tc.unfreeze, frozen_tree, is_leaf=tc.is_frozen) >>> unfrozen_tree {'a': 1, 'b': 2} """ @@ -246,7 +246,7 @@ def is_nondiff(value: Any) -> bool: False Note: - This function is meant to be used with ``jax.tree_map`` to + This function is meant to be used with ``jax.tree_util.tree_map`` to create a mask for non-differentiable nodes in a tree, that can be used to freeze the non-differentiable nodes before passing the tree to a ``jax`` transformation. @@ -338,7 +338,7 @@ def tree_mask( >>> # mask all non-differentiable nodes by default >>> def mask_if_nondiff(x): ... return tc.freeze(x) if tc.is_nondiff(x) else x - >>> masked_tree = jax.tree_map(mask_if_nondiff, tree) + >>> masked_tree = jax.tree_util.tree_map(mask_if_nondiff, tree) - Use masking on tree containing non-differentiable nodes before passing the tree to a ``jax`` transformation. @@ -411,6 +411,6 @@ def tree_unmask(tree: T, mask: MaskType = lambda _: True): >>> import jax >>> tree = [1, 2, {"a": 3, "b": 4.}] >>> # unmask all nodes - >>> tree = jax.tree_map(tc.unfreeze, tree, is_leaf=tc.is_frozen) + >>> tree = jax.tree_util.tree_map(tc.unfreeze, tree, is_leaf=tc.is_frozen) """ return _tree_mask_map(tree, mask=mask, func=unfreeze, is_leaf=is_frozen) From 5f02f834d4333a51ecac4e0f7b612b83447ece74 Mon Sep 17 00:00:00 2001 From: yspkm Date: Sun, 9 Aug 2026 09:18:04 +0800 Subject: [PATCH 3/4] ci: restore Python 3.8 JAX wheel resolution --- .github/workflows/test_jax.yml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test_jax.yml b/.github/workflows/test_jax.yml index 9a3ddb1..abdcf6c 100644 --- a/.github/workflows/test_jax.yml +++ b/.github/workflows/test_jax.yml @@ -22,7 +22,13 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install jaxlib jax + # JAX 0.4.13 is the last release that provides Python 3.8 wheels. + if [ "${{ matrix.python-version }}" = "3.8" ]; then + python -m pip install "jaxlib==0.4.13" "jax==0.4.13" \ + --find-links https://storage.googleapis.com/jax-releases/jax_releases.html + else + python -m pip install jaxlib jax + fi python -m pip install typing-extensions python -m pip install pytest pytest-benchmark wheel coverage - name: Pytest Check From 832e5ef3333c1691629c50c02915a8ae3cbea5b8 Mon Sep 17 00:00:00 2001 From: yspkm Date: Thu, 13 Aug 2026 00:02:58 +0800 Subject: [PATCH 4/4] chore(release): prepare v0.11.1 --- CHANGELOG.md | 4 ++++ pytreeclass/__init__.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a37f6f3..18a5e2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## v0.11.1 + +- Preserve inherited `TreeClass` hashability in static analysis for `@autoinit` classes. + ## v0.11.0 ## Breaking Changes: diff --git a/pytreeclass/__init__.py b/pytreeclass/__init__.py index 2135a27..fe36211 100644 --- a/pytreeclass/__init__.py +++ b/pytreeclass/__init__.py @@ -63,7 +63,7 @@ "leafwise", ) -__version__ = "0.11.0" +__version__ = "0.11.1" AtIndexer.__module__ = "pytreeclass" TreeClass.__module__ = "pytreeclass"