Merge pull request #46 from rhayes777/feature/array_impl_dict

handle converting JAX ArrayImpl to dict

autoconf/dictable.py

@@ -35,6 +35,18 @@ def nd_array_from_dict(nd_array_dict: dict, **_) -> np.ndarray:
     return np.array(nd_array_dict["array"], dtype=getattr(np, nd_array_dict["dtype"]))
 
 
+def is_array(obj) -> bool:
+    """
+    True if the object is a numpy array or an ArrayImpl (i.e. from JAX)
+    """
+    if isinstance(obj, np.ndarray):
+        return True
+    try:
+        return obj.__class__.__name__ == "ArrayImpl"
+    except AttributeError:
+        return False
+
+
 def to_dict(obj):
     if isinstance(obj, (int, float, str, bool, type(None))):
         return obj
@@ -45,7 +57,7 @@ def to_dict(obj):
         except TypeError as e:
             logger.debug(e)
 
-    if isinstance(obj, np.ndarray):
+    if is_array(obj):
         try:
             return nd_array_as_dict(obj)
         except Exception as e:

test_autoconf/test_dictable.py

@@ -16,6 +16,25 @@ def make_array():
     return np.array([1.0])
 
 
+class ArrayImpl:
+    def __init__(self, array):
+        self.array = array
+
+    @property
+    def dtype(self):
+        return self.array.dtype
+
+    def tolist(self):
+        return self.array.tolist()
+
+    def __array__(self):
+        return self.array
+
+
+def test_array_impl(array):
+    assert to_dict(ArrayImpl(array)) == to_dict(array)
+
+
 def test_array_as_dict(array_dict, array):
     assert to_dict(array) == array_dict