Fix for new jax version. (#22)

.pre-commit-config.yaml

@@ -1,17 +1,17 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
     -   id: check-merge-conflict
     -   id: debug-statements
     -   id: end-of-file-fixer
 -   repo: https://github.com/asottile/reorder_python_imports
-    rev: v3.1.0
+    rev: v3.9.0
     hooks:
     -   id: reorder-python-imports
         types: [python]
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
     -   id: check-added-large-files
         args: ['--maxkb=100']
@@ -45,12 +45,12 @@ repos:
         additional_dependencies: [black==22.3.0]
         types: [rst]
 -   repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 22.12.0
     hooks:
     -   id: black
-        language_version: python3.9
+        language_version: python3.10
 -   repo: https://github.com/PyCQA/flake8
-    rev: 4.0.1
+    rev: 5.0.4
     hooks:
     -   id: flake8
         types: [python]
@@ -71,7 +71,7 @@ repos:
             Pygments,
         ]
 -   repo: https://github.com/PyCQA/doc8
-    rev: 0.11.2
+    rev: v1.0.0
     hooks:
         -   id: doc8
 -   repo: meta
@@ -80,15 +80,15 @@ repos:
     -   id: check-useless-excludes
     # -   id: identity  # Prints all files passed to pre-commits. Debugging.
 -   repo: https://github.com/mgedmin/check-manifest
-    rev: "0.48"
+    rev: "0.49"
     hooks:
     -   id: check-manifest
 -   repo: https://github.com/PyCQA/doc8
-    rev: 0.11.2
+    rev: v1.0.0
     hooks:
     -   id: doc8
 -   repo: https://github.com/asottile/setup-cfg-fmt
-    rev: v1.20.1
+    rev: v2.2.0
     hooks:
     -   id: setup-cfg-fmt
 -   repo: https://github.com/econchick/interrogate
@@ -98,11 +98,11 @@ repos:
         args: [-v, --fail-under=20]
         exclude: ^(tests|docs|setup\.py)
 - repo: https://github.com/codespell-project/codespell
-  rev: v2.1.0
+  rev: v2.2.2
   hooks:
     -   id: codespell
 -   repo: https://github.com/asottile/pyupgrade
-    rev: v2.34.0
+    rev: v3.3.1
     hooks:
     -   id: pyupgrade
         args: [--py37-plus]

setup.cfg

@@ -17,10 +17,6 @@ classifiers =
     Operating System :: POSIX
     Programming Language :: Python :: 3
     Programming Language :: Python :: 3 :: Only
-    Programming Language :: Python :: 3.7
-    Programming Language :: Python :: 3.8
-    Programming Language :: Python :: 3.9
-    Programming Language :: Python :: 3.10
     Topic :: Scientific/Engineering
     Topic :: Utilities
 

src/pybaum/equality.py

@@ -11,10 +11,6 @@
 if IS_PANDAS_INSTALLED:
     import pandas as pd
 
-if IS_JAX_INSTALLED:
-    import jaxlib
-
-
 EQUALITY_CHECKERS = {}
 
 
@@ -28,6 +24,4 @@
 
 
 if IS_JAX_INSTALLED:
-    EQUALITY_CHECKERS[jaxlib.xla_extension.DeviceArray] = lambda a, b: bool(
-        (a == b).all()
-    )
+    EQUALITY_CHECKERS["jax.numpy.ndarray"] = lambda a, b: bool((a == b).all())

src/pybaum/registry_entries.py

@@ -1,5 +1,4 @@
 import itertools
-from collections import namedtuple
 from collections import OrderedDict
 from itertools import product
 
@@ -15,7 +14,6 @@
 
 if IS_JAX_INSTALLED:
     import jax
-    import jaxlib
 
 
 def _none():
@@ -69,7 +67,7 @@ def _tuple():
 def _namedtuple():
     """Create registry entry for namedtuple and NamedTuple."""
     entry = {
-        namedtuple: {
+        "namedtuple": {
             "flatten": lambda tree: (list(tree), tree),
             "unflatten": _unflatten_namedtuple,
             "names": lambda tree: list(tree._fields),
@@ -125,7 +123,7 @@ def _array_element_names(arr):
 def _jax_array():
     if IS_JAX_INSTALLED:
         entry = {
-            jaxlib.xla_extension.DeviceArray: {
+            "jax.numpy.ndarray": {
                 "flatten": lambda arr: (arr.flatten().tolist(), arr.shape),
                 "unflatten": lambda aux_data, leaves: jax.numpy.array(leaves).reshape(
                     aux_data

src/pybaum/typecheck.py

@@ -1,8 +1,37 @@
-from collections import namedtuple
+from pybaum.config import IS_JAX_INSTALLED
+from pybaum.config import IS_NUMPY_INSTALLED
+
+if IS_JAX_INSTALLED:
+    import jax.numpy as jnp
+
+if IS_NUMPY_INSTALLED:
+    import numpy as np
 
 
 def get_type(obj):
-    """namdetuple aware type check.
+    """Get type of candidate objects in a pytree.
+
+    This function allows us to reliably identify namedtuples, NamedTuples and jax arrays
+    for which standard ``type`` function does not work.
+
+    Args:
+        obj: The object to be checked
+
+    Returns:
+        type or str: The type of the object or a string with the type name.
+
+    """
+    if _is_namedtuple(obj):
+        out = "namedtuple"
+    elif _is_jax_array(obj):
+        out = "jax.numpy.ndarray"
+    else:
+        out = type(obj)
+    return out
+
+
+def _is_namedtuple(obj):
+    """Check if an object is a namedtuple.
 
     As in JAX we treat collections.namedtuple and typing.NamedTuple both as
     namedtuple but the exact type is preserved in the unflatten function.
@@ -24,8 +53,41 @@ def get_type(obj):
         bool
 
     """
-    if isinstance(obj, tuple) and hasattr(obj, "_fields") and hasattr(obj, "_replace"):
-        out = namedtuple
+    out = (
+        isinstance(obj, tuple) and hasattr(obj, "_fields") and hasattr(obj, "_replace")
+    )
+    return out
+
+
+def _is_jax_array(obj):
+    """Check if an object is a jax array.
+
+    The exact type of jax arrays has changed over time and is an implementation detail.
+
+    Instead we rely on isinstance checks which will likely be more stable in the future.
+    However, the behavior of isinstance for jax arrays has also changed over time. For
+    jax versions before 0.2.21, standard numpy arrays were instances of jax arrays,
+    now they are not.
+
+    Resources:
+    ----------
+
+    - https://github.com/google/jax/issues/2115
+    - https://github.com/google/jax/issues/2014
+    - https://github.com/google/jax/blob/main/CHANGELOG.md#jax-0221-sept-23-2021
+    - https://github.com/google/jax/blob/main/CHANGELOG.md#jax-0318-sep-26-2022
+
+    Args:
+        obj: The object to be checked
+
+    Returns:
+        bool
+
+    """
+    if not IS_JAX_INSTALLED:
+        out = False
+    elif IS_NUMPY_INSTALLED:
+        out = isinstance(obj, jnp.ndarray) and not isinstance(obj, np.ndarray)
     else:
-        out = type(obj)
+        out = isinstance(obj, jnp.ndarray)
     return out

tests/test_typecheck.py

@@ -6,7 +6,7 @@
 
 def test_namedtuple_is_discovered():
     bla = namedtuple("bla", ["a", "b"])(1, 2)
-    assert get_type(bla) == namedtuple
+    assert get_type(bla) == "namedtuple"
 
 
 def test_typed_namedtuple_is_discovered():
@@ -15,7 +15,7 @@ class Blubb(NamedTuple):
         b: int
 
     blubb = Blubb(1, 2)
-    assert get_type(blubb) == namedtuple
+    assert get_type(blubb) == "namedtuple"
 
 
 def test_standard_tuple_is_not_discovered():