Port InferenceData conversion code to pymc3 codebase

.github/workflows/arviz_compat.yml

@@ -7,7 +7,6 @@ on:
 
 jobs:
   pytest:
-    if: false
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest]

pymc3/__init__.py

@@ -40,7 +40,12 @@ def __set_compiler_flags():
 
 from pymc3 import gp, ode, sampling
 from pymc3.aesaraf import *
-from pymc3.backends import load_trace, save_trace
+from pymc3.backends import (
+    load_trace,
+    predictions_to_inference_data,
+    save_trace,
+    to_inference_data,
+)
 from pymc3.backends.tracetab import *
 from pymc3.blocking import *
 from pymc3.data import *

pymc3/aesaraf.py

@@ -19,10 +19,12 @@
 from aesara import scalar
 from aesara import tensor as aet
 from aesara.gradient import grad
-from aesara.graph.basic import Apply, graph_inputs
+from aesara.graph.basic import Apply, Constant, graph_inputs
 from aesara.graph.op import Op
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.tensor.elemwise import Elemwise
+from aesara.tensor.sharedvar import SharedVariable
+from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorVariable
 
 from pymc3.data import GeneratorAdapter
@@ -48,6 +50,28 @@
 ]
 
 
+def extract_obs_data(x: TensorVariable) -> np.ndarray:
+    """Extract data observed symbolic variables.
+
+    Raises
+    ------
+    TypeError
+
+    """
+    if isinstance(x, Constant):
+        return x.data
+    if isinstance(x, SharedVariable):
+        return x.get_value()
+    if x.owner and isinstance(x.owner.op, (AdvancedIncSubtensor, AdvancedIncSubtensor1)):
+        array_data = extract_obs_data(x.owner.inputs[0])
+        mask_idx = tuple(extract_obs_data(i) for i in x.owner.inputs[2:])
+        mask = np.zeros_like(array_data)
+        mask[mask_idx] = 1
+        return np.ma.MaskedArray(array_data, mask)
+
+    raise TypeError(f"Data cannot be extracted from {x}")
+
+
 def inputvars(a):
     """
     Get the inputs into a aesara variables

pymc3/backends/__init__.py

@@ -60,6 +60,7 @@
 Saved backends can be loaded using `arviz.from_netcdf`
 
 """
+from pymc3.backends.arviz import predictions_to_inference_data, to_inference_data
 from pymc3.backends.ndarray import (
     NDArray,
     load_trace,