implement index in a different way

Author: Maxim Kochurov

pymc/data.py

@@ -26,12 +26,14 @@
 import numpy as np
 
 from aesara.compile.sharedvalue import SharedVariable
+from aesara.tensor.random import RandomStream
+from aesara.tensor.random.basic import IntegersRV
 from aesara.tensor.type import TensorType
 from aesara.tensor.var import TensorConstant, TensorVariable
 
 import pymc as pm
 
-from pymc.aesaraf import at_rng, convert_observed_data
+from pymc.aesaraf import convert_observed_data
 
 __all__ = [
     "get_data",
@@ -123,6 +125,34 @@ def __hash__(self):
         return hash(id(self))
 
 
+class MinibatchIndexRV(IntegersRV):
+    _print_name = ("minibatch_index", r"\operatorname{minibatch\_index}")
+
+
+minibatch_index = MinibatchIndexRV()
+
+
+def is_minibatch(v):
+    from aesara.scalar import Cast
+    from aesara.tensor.elemwise import Elemwise
+    from aesara.tensor.subtensor import AdvancedSubtensor
+
+    return (
+        isinstance(v.owner.op, AdvancedSubtensor)
+        and isinstance(v.owner.inputs[1].owner.op, MinibatchIndexRV)
+        and (
+            v.owner.inputs[0].owner is None
+            # The only Aesara operation we allow on observed data is type casting
+            # Although we could allow for any graph that does not depend on other RVs
+            or (
+                isinstance(v.owner.inputs[0].owner.op, Elemwise)
+                and v.owner.inputs[0].owner.inputs[0].owner is None
+                and isinstance(v.owner.inputs[0].owner.op.scalar_op, Cast)
+            )
+        )
+    )
+
+
 def Minibatch(
     variable: TensorVariable, *variables: TensorVariable, batch_size: int
 ) -> Tuple[TensorVariable]:
@@ -143,17 +173,13 @@ def Minibatch(
     >>> mdata1, mdata2 = Minibatch(data1, data2, batch_size=10)
     """
 
-    def _minibatch_name(v1, v0):
-        base_name = getattr(v0, "name", "")
-        v1.name = f"minibatch_{base_name}_{id(v0)}"
-        return v1
-
-    slc = at_rng().uniform(0, variable.shape[0], size=batch_size).astype(np.int64)
+    rng = RandomStream()
+    slc = rng.gen(minibatch_index, 0, variable.shape[0], size=batch_size)
     if variables:
         variables = (variable, *variables)
-        return tuple([_minibatch_name(at.as_tensor(v)[slc], v) for v in variables])
+        return tuple([at.as_tensor(v)[slc] for v in variables])
     else:
-        return _minibatch_name(at.as_tensor(variable)[slc], variable)
+        return at.as_tensor(variable)[slc]
 
 
 def determine_coords(

pymc/distributions/logprob.py

@@ -32,6 +32,7 @@
 from aesara.tensor.var import TensorVariable
 
 from pymc.aesaraf import constant_fold, floatX
+from pymc.data import MinibatchIndexRV
 
 TOTAL_SIZE = Union[int, Sequence[int], None]
 
@@ -116,6 +117,7 @@ def _check_no_rvs(logp_terms: Sequence[TensorVariable]):
             node.owner
             and isinstance(node.owner.op, RandomVariable)
             and not isinstance(node.owner.op, SimulatorRV)
+            and not isinstance(node.owner.op, MinibatchIndexRV)
         )
     ]
     if unexpected_rv_nodes:

pymc/model.py

@@ -60,7 +60,7 @@
     replace_rvs_by_values,
 )
 from pymc.blocking import DictToArrayBijection, RaveledVars
-from pymc.data import GenTensorVariable
+from pymc.data import GenTensorVariable, is_minibatch
 from pymc.distributions.logprob import _joint_logp
 from pymc.distributions.transforms import _default_transform
 from pymc.exceptions import ImputationWarning, SamplingError, ShapeError, ShapeWarning
@@ -1311,7 +1311,7 @@ def register_rv(
                     isinstance(data.owner.op, Elemwise)
                     and isinstance(data.owner.op.scalar_op, Cast)
                 )
-                and not data.name.startswith("minibatch")
+                and not is_minibatch(data)
             ):
                 raise TypeError(
                     "Variables that depend on other nodes cannot be used for observed data."

pymc/tests/test_data.py

@@ -695,10 +695,13 @@ def test_common_errors(self):
         assert "Double Ellipsis" in str(e.value)
 
     def test_mixed1(self):
+        from pymc.distributions.logprob import joint_logp
+
         with pm.Model():
             data = np.random.rand(10, 20)
             mb = pm.Minibatch(data, batch_size=5)
-            pm.Normal("n", observed=mb, total_size=10)
+            v = pm.Normal("n", observed=mb, total_size=10)
+            assert joint_logp(v) is not None, "Check index is allowed in graph"
 
     def test_free_rv(self):
         with pm.Model() as model4: