fix(typing): set eq_default=False for autoinit

.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

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:

pytreeclass/__init__.py

@@ -63,7 +63,7 @@
     "leafwise",
 )
 
-__version__ = "0.11.0"
+__version__ = "0.11.1"
 
 AtIndexer.__module__ = "pytreeclass"
 TreeClass.__module__ = "pytreeclass"

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.
 

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)

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