Use isnull from pandas objects to capture pd.NA too

pymc-devs · Jan 26, 2021 · 6d5c78a · 6d5c78a
1 parent f0f8d83
commit 6d5c78a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymc3/model.py b/pymc3/model.py
@@ -1695,10 +1695,10 @@ 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)
+        mask = data.isnull().to_numpy()
         if mask.any():
             # there are missing values
             ret = np.ma.MaskedArray(vals, mask)