Remove non-generative incsubtensor logp inference

This is now properly done by PartiallyObservedRV

Author: ricardoV94

pymc/logprob/rewriting.py

@@ -38,8 +38,6 @@
 from collections import deque
 from collections.abc import Collection, Sequence
 
-import pytensor.tensor as pt
-
 from pytensor import config
 from pytensor.compile.mode import optdb
 from pytensor.graph.basic import (
@@ -84,7 +82,7 @@
 from pytensor.tensor.variable import TensorVariable
 
 from pymc.logprob.abstract import MeasurableVariable
-from pymc.logprob.utils import DiracDelta, indices_from_subtensor
+from pymc.logprob.utils import DiracDelta
 
 inc_subtensor_ops = (IncSubtensor, AdvancedIncSubtensor, AdvancedIncSubtensor1)
 subtensor_ops = (AdvancedSubtensor, AdvancedSubtensor1, Subtensor)
@@ -314,50 +312,6 @@ def remove_DiracDelta(fgraph, node):
     return [dd_val]
 
 
-@node_rewriter(inc_subtensor_ops)
-def incsubtensor_rv_replace(fgraph, node):
-    r"""Replace `*IncSubtensor*` `Op`\s and their value variables for log-probability calculations.
-
-    This is used to derive the log-probability graph for ``Y[idx] = data``, where
-    ``Y`` is a `RandomVariable`, ``idx`` indices, and ``data`` some arbitrary data.
-
-    To compute the log-probability of a statement like ``Y[idx] = data``, we must
-    first realize that our objective is equivalent to computing ``logprob(Y, z)``,
-    where ``z = pt.set_subtensor(y[idx], data)`` and ``y`` is the value variable
-    for ``Y``.
-
-    In other words, the log-probability for an `*IncSubtensor*` is the log-probability
-    of the underlying `RandomVariable` evaluated at ``data`` for the indices
-    given by ``idx`` and at the value variable for ``~idx``.
-
-    This provides a means of specifying "missing data", for instance.
-    """
-    rv_map_feature: PreserveRVMappings | None = getattr(fgraph, "preserve_rv_mappings", None)
-
-    if rv_map_feature is None:
-        return None  # pragma: no cover
-
-    rv_var = node.outputs[0]
-    if rv_var not in rv_map_feature.rv_values:
-        return None  # pragma: no cover
-
-    base_rv_var = node.inputs[0]
-
-    if not rv_map_feature.request_measurable([base_rv_var]):
-        return None
-
-    data = node.inputs[1]
-    idx = indices_from_subtensor(getattr(node.op, "idx_list", None), node.inputs[2:])
-
-    # Create a new value variable with the indices `idx` set to `data`
-    value_var = rv_map_feature.rv_values[rv_var]
-    new_value_var = pt.set_subtensor(value_var[idx], data)
-    rv_map_feature.update_rv_maps(rv_var, new_value_var, base_rv_var)
-
-    # Return the `RandomVariable` being indexed
-    return [base_rv_var]
-
-
 logprob_rewrites_db = SequenceDB()
 logprob_rewrites_db.name = "logprob_rewrites_db"
 # Introduce sigmoid. We do it before canonicalization so that useless mul are removed next
@@ -378,7 +332,6 @@ def incsubtensor_rv_replace(fgraph, node):
 # (or eventually) the graph outputs.  Often this is done by lifting other `Op`s
 # "up" through the random/measurable variables and into their inputs.
 measurable_ir_rewrites_db.register("subtensor_lift", local_subtensor_rv_lift, "basic")
-measurable_ir_rewrites_db.register("incsubtensor_lift", incsubtensor_rv_replace, "basic")
 
 # These rewrites are used to introduce specalized operations with better logprob graphs
 specialization_ir_rewrites_db = EquilibriumDB()

tests/logprob/test_basic.py

@@ -44,11 +44,6 @@
 
 from pytensor.graph.basic import ancestors, equal_computations
 from pytensor.tensor.random.op import RandomVariable
-from pytensor.tensor.subtensor import (
-    AdvancedIncSubtensor,
-    AdvancedIncSubtensor1,
-    IncSubtensor,
-)
 
 import pymc as pm
 
@@ -173,20 +168,6 @@ def test_factorized_joint_logprob_diff_dims():
     assert exp_logp_val == pytest.approx(logp_val)
 
 
-def test_incsubtensor_original_values_output_dict():
-    """
-    Test that the original un-incsubtensor value variable appears an the key of
-    the logprob factor
-    """
-
-    base_rv = pt.random.normal(0, 1, size=2)
-    rv = pt.set_subtensor(base_rv[0], 5)
-    vv = rv.clone()
-
-    logp_dict = conditional_logp({rv: vv})
-    assert vv in logp_dict
-
-
 def test_persist_inputs():
     """Make sure we don't unnecessarily clone variables."""
     x = pt.scalar("x")
@@ -276,54 +257,6 @@ def test_joint_logp_basic():
     assert a_value_var in res_ancestors
 
 
-@pytest.mark.parametrize(
-    "indices, size",
-    [
-        (slice(0, 2), 5),
-        (np.r_[True, True, False, False, True], 5),
-        (np.r_[0, 1, 4], 5),
-        ((np.array([0, 1, 4]), np.array([0, 1, 4])), (5, 5)),
-    ],
-)
-def test_joint_logp_incsubtensor(indices, size):
-    """Make sure we can compute a log-likelihood for ``Y[idx] = data`` where ``Y`` is univariate."""
-
-    mu = pm.floatX(np.power(10, np.arange(np.prod(size)))).reshape(size)
-    data = mu[indices]
-    sigma = 0.001
-    rng = np.random.default_rng(232)
-    a_val = rng.normal(mu, sigma, size=size).astype(pytensor.config.floatX)
-
-    rng = pytensor.shared(rng, borrow=False)
-    a = pm.Normal.dist(mu, sigma, size=size, rng=rng)
-    a_value_var = a.type()
-    a.name = "a"
-
-    a_idx = pt.set_subtensor(a[indices], data)
-
-    assert isinstance(a_idx.owner.op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1)
-
-    a_idx_value_var = a_idx.type()
-    a_idx_value_var.name = "a_idx_value"
-
-    a_idx_logp = transformed_conditional_logp(
-        (a_idx,),
-        rvs_to_values={a_idx: a_value_var},
-        rvs_to_transforms={},
-    )
-
-    logp_vals = a_idx_logp[0].eval({a_value_var: a_val})
-
-    # The indices that were set should all have the same log-likelihood values,
-    # because the values they were set to correspond to the unique means along
-    # that dimension.  This helps us confirm that the log-likelihood is
-    # associating the assigned values with their correct parameters.
-    a_val_idx = a_val.copy()
-    a_val_idx[indices] = data
-    exp_obs_logps = sp.norm.logpdf(a_val_idx, mu, sigma)
-    np.testing.assert_almost_equal(logp_vals, exp_obs_logps)
-
-
 def test_model_unchanged_logprob_access():
     # Issue #5007
     with pm.Model() as model:

tests/logprob/test_rewriting.py

@@ -34,19 +34,13 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 
-import numpy as np
 import pytensor.tensor as pt
-import pytest
-import scipy.stats.distributions as sp
 
 from pytensor.graph import ancestors
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.subtensor import (
-    AdvancedIncSubtensor,
-    AdvancedIncSubtensor1,
-    IncSubtensor,
     Subtensor,
 )
 
@@ -105,41 +99,3 @@ def test_local_remove_TransformedVariable():
     [p_logp] = conditional_logp({p_rv: p_vv}, extra_rewrites=tr).values()
 
     assert not any(isinstance(v.owner.op, TransformedValue) for v in ancestors([p_logp]) if v.owner)
-
-
-@pytest.mark.parametrize(
-    "indices, size",
-    [
-        (slice(0, 2), 5),
-        (np.r_[True, True, False, False, True], 5),
-        (np.r_[0, 1, 4], 5),
-        ((np.array([0, 1, 4]), np.array([0, 1, 4])), (5, 5)),
-    ],
-)
-def test_joint_logprob_incsubtensor(indices, size):
-    """Make sure we can compute a joint log-probability for ``Y[idx] = data`` where ``Y`` is univariate."""
-
-    rng = np.random.RandomState(232)
-    mu = np.power(10, np.arange(np.prod(size))).reshape(size)
-    sigma = 0.001
-    data = rng.normal(mu[indices], 1.0)
-    y_val = rng.normal(mu, sigma, size=size)
-
-    Y_base_rv = pt.random.normal(mu, sigma, size=size)
-    Y_rv = pt.set_subtensor(Y_base_rv[indices], data)
-    Y_rv.name = "Y"
-    y_value_var = Y_rv.clone()
-    y_value_var.name = "y"
-
-    assert isinstance(Y_rv.owner.op, IncSubtensor | AdvancedIncSubtensor | AdvancedIncSubtensor1)
-
-    Y_rv_logp = conditional_logp({Y_rv: y_value_var})
-    Y_rv_logp_combined = pt.add(*Y_rv_logp.values())
-
-    obs_logps = Y_rv_logp_combined.eval({y_value_var: y_val})
-
-    y_val_idx = y_val.copy()
-    y_val_idx[indices] = data
-    exp_obs_logps = sp.norm.logpdf(y_val_idx, mu, sigma)
-
-    np.testing.assert_almost_equal(obs_logps, exp_obs_logps)