Handles objects with __array__ (For jax.numpy.ndarray support) (#2481)

* handles objects with __array__

* format

* added __array__ into docstring

* black formatting

---------

Co-authored-by: clarkmiyamoto <[email protected]>

Author: clarkmiyamoto

arviz/data/converters.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import xarray as xr
+import pandas as pd
 
 try:
     from tree import is_nested
@@ -44,6 +45,8 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
             | dict: creates an xarray dataset as the only group
             | numpy array: creates an xarray dataset as the only group, gives the
                          array an arbitrary name
+            | object with __array__: converts to numpy array, then creates an xarray dataset as
+                         the only group, gives the array an arbitrary name
     group : str
         If `obj` is a dict or numpy array, assigns the resulting xarray
         dataset to this group. Default: "posterior".
@@ -115,6 +118,13 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
         dataset = dict_to_dataset(obj, coords=coords, dims=dims)
     elif isinstance(obj, np.ndarray):
         dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
+    elif (
+        hasattr(obj, "__array__")
+        and callable(getattr(obj, "__array__"))
+        and (not isinstance(obj, pd.DataFrame))
+    ):
+        obj = obj.__array__()
+        dataset = dict_to_dataset({"x": obj}, coords=coords, dims=dims)
     elif isinstance(obj, (list, tuple)) and isinstance(obj[0], str) and obj[0].endswith(".csv"):
         if group == "sample_stats":
             kwargs["posterior"] = kwargs.pop(group)
@@ -129,6 +139,7 @@ def convert_to_inference_data(obj, *, group="posterior", coords=None, dims=None,
             "pytree (if 'dm-tree' is installed)",
             "netcdf filename",
             "numpy array",
+            "object with __array__",
             "pystan fit",
             "emcee fit",
             "pyro mcmc fit",

arviz/data/io_pyjags.py

@@ -277,7 +277,7 @@ def _extract_arviz_dict_from_inference_data(
 
 
 def _convert_arviz_dict_to_pyjags_dict(
-    samples: tp.Mapping[str, np.ndarray]
+    samples: tp.Mapping[str, np.ndarray],
 ) -> tp.Mapping[str, np.ndarray]:
     """
     Convert and ArviZ dictionary to a PyJAGS dictionary.