Avoid cloning of Minibatch values

Author: ricardoV94

pymc/data.py

@@ -36,6 +36,7 @@
 
 import pymc as pm
 
+from pymc.logprob.abstract import _get_measurable_outputs
 from pymc.pytensorf import convert_observed_data
 
 __all__ = [
@@ -134,6 +135,11 @@ def make_node(self, rng, *args, **kwargs):
         return super().make_node(rng, *args, **kwargs)
 
 
+@_get_measurable_outputs.register(MinibatchIndexRV)
+def minibatch_index_rv_measuarable_outputs(op, node):
+    return []
+
+
 minibatch_index = MinibatchIndexRV()
 
 

pymc/logprob/basic.py

@@ -44,7 +44,13 @@
 import pytensor.tensor as pt
 
 from pytensor import config
-from pytensor.graph.basic import Variable, graph_inputs, io_toposort
+from pytensor.graph.basic import (
+    Constant,
+    Variable,
+    ancestors,
+    graph_inputs,
+    io_toposort,
+)
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.rewriting.basic import GraphRewriter, NodeRewriter
 from pytensor.tensor.random.op import RandomVariable
@@ -231,10 +237,16 @@ def factorized_joint_logprob(
     # node.
     replacements = updated_rv_values.copy()
 
-    # To avoid cloning the value variables, we map them to themselves in the
-    # `replacements` `dict` (i.e. entries already existing in `replacements`
-    # aren't cloned)
-    replacements.update({v: v for v in rv_values.values()})
+    # To avoid cloning the value variables (or ancestors of value variables),
+    # we map them to themselves in the `replacements` `dict`
+    # (i.e. entries already existing in `replacements` aren't cloned)
+    replacements.update(
+        {
+            v: v
+            for v in ancestors(rv_values.values())
+            if (not isinstance(v, Constant) and v not in replacements)
+        }
+    )
 
     # Walk the graph from its inputs to its outputs and construct the
     # log-probability

pymc/logprob/rewriting.py

@@ -39,7 +39,7 @@
 import pytensor.tensor as pt
 
 from pytensor.compile.mode import optdb
-from pytensor.graph.basic import Variable
+from pytensor.graph.basic import Constant, Variable, ancestors
 from pytensor.graph.features import Feature
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import GraphRewriter, node_rewriter
@@ -316,8 +316,8 @@ def construct_ir_fgraph(
     # the old nodes to the new ones; otherwise, we won't be able to use
     # `rv_values`.
     # We start the `dict` with mappings from the value variables to themselves,
-    # to prevent them from being cloned.
-    memo = {v: v for v in rv_values.values()}
+    # to prevent them from being cloned. This also includes ancestors
+    memo = {v: v for v in ancestors(rv_values.values()) if not isinstance(v, Constant)}
 
     # We add `ShapeFeature` because it will get rid of references to the old
     # `RandomVariable`s that have been lifted; otherwise, it will be difficult

tests/variational/test_minibatch_rv.py

@@ -20,6 +20,7 @@
 import pymc as pm
 
 from pymc import Normal, draw
+from pymc.data import minibatch_index
 from pymc.testing import select_by_precision
 from pymc.variational.minibatch_rv import create_minibatch_rv
 from tests.test_data import gen1, gen2
@@ -155,3 +156,36 @@ def test_random(self):
         mx = create_minibatch_rv(x, total_size=(10,))
         assert mx is not x
         np.testing.assert_array_equal(draw(mx, random_seed=1), draw(x, random_seed=1))
+
+    @pytest.mark.filterwarnings("error")
+    def test_minibatch_parameter_and_value(self):
+        rng = np.random.default_rng(161)
+        total_size = 1000
+
+        with pm.Model(check_bounds=False) as m:
+            AD = pm.MutableData("AD", np.arange(total_size, dtype="float64"))
+            TD = pm.MutableData("TD", np.arange(total_size, dtype="float64"))
+
+            minibatch_idx = minibatch_index(0, 10, size=(9,))
+            AD_mt = AD[minibatch_idx]
+            TD_mt = TD[minibatch_idx]
+
+            pm.Normal(
+                "AD_predicted",
+                mu=TD_mt,
+                observed=AD_mt,
+                total_size=1000,
+            )
+
+        logp_fn = m.compile_logp()
+
+        ip = m.initial_point()
+        np.testing.assert_allclose(logp_fn(ip), st.norm.logpdf(0) * 1000)
+
+        with m:
+            pm.set_data({"AD": np.arange(total_size) + 1})
+        np.testing.assert_allclose(logp_fn(ip), st.norm.logpdf(1) * 1000)
+
+        with m:
+            pm.set_data({"AD": rng.normal(size=1000)})
+        assert logp_fn(ip) != logp_fn(ip)