Fix Minibatch for multiple variables

Author: ricardoV94

pymc/data.py

@@ -170,13 +170,11 @@ def assert_all_scalars_equal(scalar, *scalars):
     else:
         return Assert(
             "All variables shape[0] in Minibatch should be equal, check your Minibatch(data1, data2, ...) code"
-        )(scalar, pt.all([scalar == s for s in scalars]))
+        )(scalar, pt.all([pt.eq(scalar, s) for s in scalars]))
 
 
 def Minibatch(variable: TensorVariable, *variables: TensorVariable, batch_size: int):
-    """
-    Get random slices from variables from the leading dimension.
-
+    """Get random slices from variables from the leading dimension.
 
     Parameters
     ----------
@@ -191,6 +189,9 @@ def Minibatch(variable: TensorVariable, *variables: TensorVariable, batch_size:
     >>> mdata1, mdata2 = Minibatch(data1, data2, batch_size=10)
     """
 
+    if not isinstance(batch_size, int):
+        raise TypeError("batch_size must be an integer")
+
     tensor, *tensors = tuple(map(pt.as_tensor, (variable, *variables)))
     upper = assert_all_scalars_equal(*[t.shape[0] for t in (tensor, *tensors)])
     slc = minibatch_index(0, upper, size=batch_size)

tests/test_data.py

@@ -14,6 +14,7 @@
 
 import io
 import itertools as it
+import re
 
 import cloudpickle
 import numpy as np
@@ -614,3 +615,23 @@ def test_assert(self):
         ):
             d1, d2 = pm.Minibatch(self.data, self.data[::2], batch_size=20)
             d1.eval()
+
+    def test_multiple_vars(self):
+        A = np.arange(1000)
+        B = np.arange(1000)
+        mA, mB = pm.Minibatch(A, B, batch_size=10)
+
+        [draw_mA, draw_mB] = pm.draw([mA, mB])
+        assert draw_mA.shape == (10,)
+        np.testing.assert_allclose(draw_mA, draw_mB)
+
+        # Check invalid dims
+        A = np.arange(1000)
+        C = np.arange(999)
+        mA, mC = pm.Minibatch(A, C, batch_size=10)
+
+        with pytest.raises(
+            AssertionError,
+            match=re.escape("All variables shape[0] in Minibatch should be equal"),
+        ):
+            pm.draw([mA, mC])