Use isnull from pandas objects to capture pd.NA too

pymc-devs · Jan 25, 2021 · 8476aa4 · 8476aa4
1 parent f0f8d83
commit 8476aa4
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pymc3/model.py b/pymc3/model.py
@@ -1695,11 +1695,11 @@ def pandas_to_array(data):
     XXX: When `data` is a generator, this will return a Theano tensor!
 
     """
-    if hasattr(data, "to_numpy"):
+    if hasattr(data, "to_numpy") and hasattr(data, "isnull"):
         # typically, but not limited to pandas objects
         vals = data.to_numpy()
-        mask = np.isnan(vals)
-        if mask.any():
+        mask = data.isnull().to_numpy()
+        if mask.any().any():
             # there are missing values
             ret = np.ma.MaskedArray(vals, mask)
         else: