Add _as_tensor_variable converter for pandas objects

pymc/aesaraf.py

@@ -50,6 +50,7 @@
 from aesara.graph.op import Op, compute_test_value
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.scalar.basic import Cast
+from aesara.tensor.basic import _as_tensor_variable
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import (
@@ -142,6 +143,20 @@ def convert_observed_data(data):
         return floatX(ret)
 
 
+try:
+    import pandas as pd
+
+    _pd_available: bool = True
+except ModuleNotFoundError:
+    _pd_available: bool = False
+
+if _pd_available:
+
+    @_as_tensor_variable.register(pd.DataFrame)
+    def dataframe_to_tensor_variable(df: pd.DataFrame, *args, **kwargs) -> TensorVariable:
+        return at.as_tensor_variable(df.to_numpy(), *args, **kwargs)
+
+
 def change_rv_size(
     rv: TensorVariable,
     new_size: PotentialShapeType,

pymc/tests/test_aesaraf.py

@@ -47,6 +47,29 @@
 from pymc.vartypes import int_types
 
 
+@pytest.mark.parametrize(
+    argnames="np_array",
+    argvalues=[np.array([[1.0], [2.0], [-1.0]]), np.array([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]])],
+)
+def test_pd_as_tensor_variable(np_array: np.ndarray) -> None:
+    pd = pytest.importorskip("pandas")
+    df: pd.DataFrame = pd.DataFrame(np_array)
+    np.testing.assert_array_equal(x=at.as_tensor_variable(x=df).eval(), y=np_array)
+
+
+def test_pd_as_tensor_variable_multiindex() -> None:
+    pd = pytest.importorskip("pandas")
+
+    tuples = [("L", "Q"), ("L", "I"), ("O", "L"), ("O", "I")]
+
+    index = pd.MultiIndex.from_tuples(tuples, names=["Id1", "Id2"])
+
+    df = pd.DataFrame({"A": [12.0, 80.0, 30.0, 20.0], "B": [120.0, 700.0, 30.0, 20.0]}, index=index)
+    np_array = np.array([[12.0, 80.0, 30.0, 20.0], [120.0, 700.0, 30.0, 20.0]]).T
+    assert isinstance(df.index, pd.MultiIndex)
+    np.testing.assert_array_equal(x=at.as_tensor_variable(x=df).eval(), y=np_array)
+
+
 def test_change_rv_size():
     loc = at.as_tensor_variable([1, 2])
     rv = normal(loc=loc)