Allow registering XTensorVariables directly in model

Author: ricardoV94

pymc/dims/__init__.py

@@ -39,6 +39,9 @@ def __init__():
     logprob_rewrites_db.register(
         "pre_lower_xtensor", optdb.query("+lower_xtensor"), "basic", position=0.1
     )
+    logprob_rewrites_db.register(
+        "post_lower_xtensor", optdb.query("+lower_xtensor"), "cleanup", position=5.1
+    )
     initial_point_rewrites_db.register(
         "lower_xtensor", optdb.query("+lower_xtensor"), "basic", position=0.1
     )

pymc/dims/distributions/core.py

@@ -13,18 +13,26 @@
 #   limitations under the License.
 from collections.abc import Callable, Sequence
 from itertools import chain
+from typing import cast
 
+import numpy as np
+
+from pytensor.graph import node_rewriter
 from pytensor.graph.basic import Variable
 from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.random.op import RandomVariable
 from pytensor.xtensor import as_xtensor
+from pytensor.xtensor.basic import XTensorFromTensor, xtensor_from_tensor
 from pytensor.xtensor.type import XTensorVariable
 
-from pymc import modelcontext
-from pymc.dims.model import with_dims
-from pymc.distributions import transforms
+from pymc import SymbolicRandomVariable, modelcontext
+from pymc.dims.distributions.transforms import DimTransform, log_odds_transform, log_transform
 from pymc.distributions.distribution import _support_point, support_point
 from pymc.distributions.shape_utils import DimsWithEllipsis, convert_dims_with_ellipsis
-from pymc.logprob.transforms import Transform
+from pymc.logprob.abstract import MeasurableOp, _logprob
+from pymc.logprob.rewriting import measurable_ir_rewrites_db
+from pymc.logprob.tensor import MeasurableDimShuffle
+from pymc.logprob.utils import filter_measurable_variables
 from pymc.util import UNSET
 
 
@@ -36,25 +44,98 @@ def dimshuffle_support_point(ds_op, _, rv):
     return ds_op(support_point(rv))
 
 
+@_support_point.register(XTensorFromTensor)
+def xtensor_from_tensor_support_point(xtensor_op, _, rv):
+    # We remove the xtensor_from_tensor operation, so initial_point doesn't have to do a further lowering
+    return xtensor_op(support_point(rv))
+
+
+class MeasurableXTensorFromTensor(MeasurableOp, XTensorFromTensor):
+    __props__ = ("dims", "core_dims")  # type: ignore[assignment]
+
+    def __init__(self, dims, core_dims):
+        super().__init__(dims=dims)
+        self.core_dims = tuple(core_dims) if core_dims is not None else None
+
+
+@node_rewriter([XTensorFromTensor])
+def find_measurable_xtensor_from_tensor(fgraph, node) -> list[XTensorVariable] | None:
+    if isinstance(node.op, MeasurableXTensorFromTensor):
+        return None
+
+    xs = filter_measurable_variables(node.inputs)
+
+    if not xs:
+        # Check if we have a transposition instead
+        # The rewrite that introduces measurable tranpsoses refuses to apply to multivariate RVs
+        # So we have a chance of inferring the core dims!
+        [ds] = node.inputs
+        ds_node = ds.owner
+        if not (
+            ds_node is not None
+            and isinstance(ds_node.op, DimShuffle)
+            and ds_node.op.is_transpose
+            and filter_measurable_variables(ds_node.inputs)
+        ):
+            return None
+        [x] = ds_node.inputs
+        if not (
+            x.owner is not None and isinstance(x.owner.op, RandomVariable | SymbolicRandomVariable)
+        ):
+            return None
+
+        measurable_x = MeasurableDimShuffle(**ds_node.op._props_dict())(x)  # type: ignore[attr-defined]
+
+        ndim_supp = x.owner.op.ndim_supp
+        if ndim_supp:
+            inverse_transpose = np.argsort(ds_node.op.shuffle)
+            dims = node.op.dims
+            dims_before_transpose = tuple(dims[i] for i in inverse_transpose)
+            core_dims = dims_before_transpose[-ndim_supp:]
+        else:
+            core_dims = ()
+
+        new_out = MeasurableXTensorFromTensor(dims=node.op.dims, core_dims=core_dims)(measurable_x)
+    else:
+        # If this happens we know there's no measurable transpose in between and we can
+        # safely infer the core_dims positionally when the inner logp is returned
+        new_out = MeasurableXTensorFromTensor(dims=node.op.dims, core_dims=None)(*node.inputs)
+    return [cast(XTensorVariable, new_out)]
+
+
+@_logprob.register(MeasurableXTensorFromTensor)
+def measurable_xtensor_from_tensor(op, values, rv, **kwargs):
+    rv_logp = _logprob(rv.owner.op, tuple(v.values for v in values), *rv.owner.inputs, **kwargs)
+    if op.core_dims is None:
+        # The core_dims of the inner rv are on the right
+        dims = op.dims[: rv_logp.ndim]
+    else:
+        # We inferred where the core_dims are!
+        dims = [d for d in op.dims if d not in op.core_dims]
+    return xtensor_from_tensor(rv_logp, dims=dims)
+
+
+measurable_ir_rewrites_db.register(
+    "measurable_xtensor_from_tensor", find_measurable_xtensor_from_tensor, "basic", "xtensor"
+)
+
+
 class DimDistribution:
     """Base class for PyMC distribution that wrap pytensor.xtensor.random operations, and follow xarray-like semantics."""
 
     xrv_op: Callable
-    default_transform: Transform | None = None
+    default_transform: DimTransform | None = None
 
     @staticmethod
     def _as_xtensor(x):
         try:
             return as_xtensor(x)
         except TypeError:
-            try:
-                return with_dims(x)
-            except ValueError:
-                raise ValueError(
-                    f"Variable {x} must have dims associated with it.\n"
-                    "To avoid subtle bugs, PyMC does not make any assumptions about the dims of parameters.\n"
-                    "Use `as_xtensor` with the `dims` keyword argument to specify the dims explicitly."
-                )
+            raise ValueError(
+                f"Variable {x} must have dims associated with it.\n"
+                "To avoid subtle bugs, PyMC does not make any assumptions about the dims of parameters.\n"
+                "Use `as_xtensor` with the `dims` keyword argument to specify the dims explicitly."
+            )
 
     def __new__(
         cls,
@@ -119,10 +200,22 @@ def __new__(
         else:
             # Align observed dims with those of the RV
             # TODO: If this fails give a more informative error message
-            observed = observed.transpose(*rv_dims).values
+            observed = observed.transpose(*rv_dims)
+
+        # Check user didn't pass regular transforms
+        if transform not in (UNSET, None):
+            if not isinstance(transform, DimTransform):
+                raise TypeError(
+                    f"Transform must be a DimTransform, form pymc.dims.transforms, but got {type(transform)}."
+                )
+        if default_transform not in (UNSET, None):
+            if not isinstance(default_transform, DimTransform):
+                raise TypeError(
+                    f"default_transform must be a DimTransform, from pymc.dims.transforms, but got {type(default_transform)}."
+                )
 
         rv = model.register_rv(
-            rv.values,
+            rv,
             name=name,
             observed=observed,
             total_size=total_size,
@@ -182,10 +275,10 @@ def dist(self, *args, core_dims: str | Sequence[str] | None = None, **kwargs):
 class PositiveDimDistribution(DimDistribution):
     """Base class for positive continuous distributions."""
 
-    default_transform = transforms.log
+    default_transform = log_transform
 
 
 class UnitDimDistribution(DimDistribution):
     """Base class for unit-valued distributions."""
 
-    default_transform = transforms.logodds
+    default_transform = log_odds_transform

pymc/dims/distributions/transforms.py

@@ -0,0 +1,53 @@
+#   Copyright 2025 - present The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   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 pytensor.xtensor as ptx
+
+from pymc.logprob.transforms import Transform
+
+
+class DimTransform(Transform):
+    """Base class for transforms that are applied to dim distriubtions."""
+
+
+class LogTransform(DimTransform):
+    name = "log"
+
+    def forward(self, value, *inputs):
+        return ptx.math.log(value)
+
+    def backward(self, value, *inputs):
+        return ptx.math.exp(value)
+
+    def log_jac_det(self, value, *inputs):
+        return value
+
+
+log_transform = LogTransform()
+
+
+class LogOddsTransform(DimTransform):
+    name = "logodds"
+
+    def backward(self, value, *inputs):
+        return ptx.math.expit(value)
+
+    def forward(self, value, *inputs):
+        return ptx.math.log(value / (1 - value))
+
+    def log_jac_det(self, value, *inputs):
+        sigmoid_value = ptx.math.sigmoid(value)
+        return ptx.math.log(sigmoid_value) + ptx.math.log1p(-sigmoid_value)
+
+
+log_odds_transform = LogOddsTransform()

pymc/dims/model.py

@@ -15,7 +15,6 @@
 
 from pytensor.tensor import TensorVariable
 from pytensor.xtensor import as_xtensor
-from pytensor.xtensor.basic import TensorFromXTensor
 from pytensor.xtensor.type import XTensorVariable
 
 from pymc.data import Data as RegularData
@@ -30,38 +29,13 @@
 from pymc.model.core import Potential as RegularPotential
 
 
-def with_dims(x: TensorVariable | XTensorVariable, model: Model | None = None) -> XTensorVariable:
-    """Recover the dims of a variable that was registered in the Model."""
-    if isinstance(x, XTensorVariable):
-        return x
-
-    if (x.owner is not None) and isinstance(x.owner.op, TensorFromXTensor):
-        dims = x.owner.inputs[0].type.dims
-        return as_xtensor(x, dims=dims, name=x.name)
-
-    # Try accessing the model context to get dims
-    try:
-        model = modelcontext(model)
-        if (
-            model.named_vars.get(x.name, None) is x
-            and (dims := model.named_vars_to_dims.get(x.name, None)) is not None
-        ):
-            return as_xtensor(x, dims=dims, name=x.name)
-    except TypeError:
-        pass
-
-    raise ValueError(f"variable {x} doesn't have dims associated with it")
-
-
 def Data(
     name: str, value, dims: Dims = None, model: Model | None = None, **kwargs
 ) -> XTensorVariable:
     """Wrapper around pymc.Data that returns an XtensorVariable.
 
     Dimensions are required if the input is not a scalar.
     These are always forwarded to the model object.
-
-    The respective TensorVariable is registered in the model
     """
     model = modelcontext(model)
     dims = convert_dims(dims)  # type: ignore[assignment]
@@ -90,12 +64,9 @@ def _register_and_return_xtensor_variable(
             value = value.transpose(*dims)
         # Regardless of whether dims are provided, we now have them
         dims = value.type.dims
-        # Register the equivalent TensorVariable with the model so it doesn't see XTensorVariables directly.
-        value = value.values  # type: ignore[union-attr]
-
-    value = registration_func(name, value, dims=dims, model=model)
-
-    return as_xtensor(value, dims=dims, name=name)  # type: ignore[arg-type]
+    else:
+        value = as_xtensor(value, dims=dims, name=name)  # type: ignore[arg-type]
+    return registration_func(name, value, dims=dims, model=model)
 
 
 def Deterministic(
@@ -107,8 +78,6 @@ def Deterministic(
     If the input is not an XTensorVariable, it is converted to one using `as_xtensor`. Dims are required if the input is not a scalar.
 
     The dimensions of the resulting XTensorVariable are always forwarded to the model object.
-
-    The respective TensorVariable is registered in the model
     """
     return _register_and_return_xtensor_variable(name, value, dims, model, RegularDeterministic)
 
@@ -122,7 +91,5 @@ def Potential(
     If the input is not an XTensorVariable, it is converted to one using `as_xtensor`. Dims are required if the input is not a scalar.
 
     The dimensions of the resulting XTensorVariable are always forwarded to the model object.
-
-    The respective TensorVariable is registered in the model.
     """
     return _register_and_return_xtensor_variable(name, value, dims, model, RegularPotential)