compile_pymc handles default updates for duplicated nodes

Fixes bug in VI with multiple Minibatch variables, which occurred due to separate calls to model.logp (from model.datalogp and model.varlogp) that create distinct clones of the RandomIntegersRV underlying minibatch slicing. `compile_pymc` would not set any updates in this case

Author: ricardoV94

pymc/pytensorf.py

@@ -32,6 +32,7 @@
     Constant,
     Variable,
     clone_get_equiv,
+    equal_computations,
     graph_inputs,
     walk,
 )
@@ -872,8 +873,23 @@ def find_default_update(clients, rng: Variable) -> None | Variable:
             return rng
 
         if len(rng_clients) > 1:
+            # Multiple clients are techincally fine if they are used in identical operations
+            # We check if the default_update of each client would be the same
+            update, *other_updates = (
+                find_default_update(
+                    # Pass version of clients that includes only one the RNG clients at a time
+                    clients | {rng: [rng_client]},
+                    rng,
+                )
+                for rng_client in rng_clients
+            )
+            if all(equal_computations([update], [other_update]) for other_update in other_updates):
+                return update
+
             warnings.warn(
-                f"RNG Variable {rng} has multiple clients. This is likely an inconsistent random graph.",
+                f"RNG Variable {rng} has multiple distinct clients {rng_clients}, "
+                f"likely due to an inconsistent random graph. "
+                f"No default update will be returned.",
                 UserWarning,
             )
             return None

tests/test_pytensorf.py

@@ -11,7 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import warnings
 
 import numpy as np
 import numpy.ma as ma
@@ -486,34 +485,45 @@ def test_random_seed(self):
         assert x3_eval == x2_eval
         assert y3_eval == y2_eval
 
+    @pytest.mark.filterwarnings("error")  # This is part of the test
     def test_multiple_updates_same_variable(self):
-        # Raise if unexpected warning is issued
-        with warnings.catch_warnings():
-            warnings.simplefilter("error")
-
-            rng = pytensor.shared(np.random.default_rng(), name="rng")
-            x = pt.random.normal(rng=rng)
-            y = pt.random.normal(rng=rng)
-
-            # No warnings if only one variable is used
-            assert compile_pymc([], [x])
-            assert compile_pymc([], [y])
-
-            user_warn_msg = "RNG Variable rng has multiple clients"
-            with pytest.warns(UserWarning, match=user_warn_msg):
-                f = compile_pymc([], [x, y], random_seed=456)
-            assert f() == f()
-
-            # The user can provide an explicit update, but we will still issue a warning
-            with pytest.warns(UserWarning, match=user_warn_msg):
-                f = compile_pymc([], [x, y], updates={rng: y.owner.outputs[0]}, random_seed=456)
-            assert f() != f()
-
-            # Same with default update
-            rng.default_update = x.owner.outputs[0]
-            with pytest.warns(UserWarning, match=user_warn_msg):
-                f = compile_pymc([], [x, y], updates={rng: y.owner.outputs[0]}, random_seed=456)
-            assert f() != f()
+        rng = pytensor.shared(np.random.default_rng(), name="rng")
+        x = pt.random.normal(0, rng=rng)
+        y = pt.random.normal(1, rng=rng)
+
+        # No warnings if only one variable is used
+        assert compile_pymc([], [x])
+        assert compile_pymc([], [y])
+
+        user_warn_msg = "RNG Variable rng has multiple distinct clients"
+        with pytest.warns(UserWarning, match=user_warn_msg):
+            f = compile_pymc([], [x, y], random_seed=456)
+        assert f() == f()
+
+        # The user can provide an explicit update, but we will still issue a warning
+        with pytest.warns(UserWarning, match=user_warn_msg):
+            f = compile_pymc([], [x, y], updates={rng: y.owner.outputs[0]}, random_seed=456)
+        assert f() != f()
+
+        # Same with default update
+        rng.default_update = x.owner.outputs[0]
+        with pytest.warns(UserWarning, match=user_warn_msg):
+            f = compile_pymc([], [x, y], updates={rng: y.owner.outputs[0]}, random_seed=456)
+        assert f() != f()
+
+    @pytest.mark.filterwarnings("error")  # This is part of the test
+    def test_duplicated_client_nodes(self):
+        """Test compile_pymc can handle duplicated (mergeable) RV updates."""
+        rng = pytensor.shared(np.random.default_rng(1))
+        x = pt.random.normal(rng=rng)
+        y = x.owner.clone().default_output()
+
+        fn = compile_pymc([], [x, y], random_seed=1)
+        res_x1, res_y1 = fn()
+        assert res_x1 == res_y1
+        res_x2, res_y2 = fn()
+        assert res_x2 == res_y2
+        assert res_x1 != res_x2
 
     def test_nested_updates(self):
         rng = pytensor.shared(np.random.default_rng())

tests/variational/test_inference.py

@@ -472,3 +472,25 @@ def test_fit_data_coords(hierarchical_model, hierarchical_model_data):
             hierarchical_model_data["group_coords"].keys()
         )
         assert data["mu"].shape == tuple()
+
+
+def test_multiple_minibatch_variables():
+    """Regression test for bug reported in
+    https://discourse.pymc.io/t/verifying-that-minibatch-is-actually-randomly-sampling/14308
+    """
+    true_weights = np.array([-5, 5] * 5)
+    feature = np.repeat(np.eye(10), 10_000, axis=0)
+    y = feature @ true_weights
+
+    with pm.Model() as model:
+        minibatch_feature, minibatch_y = pm.Minibatch(feature, y, batch_size=1)
+        weights = pm.Normal("weights", 0, 10, shape=10)
+        pm.Normal(
+            "y",
+            mu=minibatch_feature @ weights,
+            sigma=0.01,
+            observed=minibatch_y,
+            total_size=len(y),
+        )
+        mean_field = pm.fit(10_000, obj_optimizer=pm.adam(learning_rate=0.01), progressbar=False)
+    np.testing.assert_allclose(mean_field.mean.get_value(), true_weights, rtol=1e-1)