Allow registering XTensorVariables directly in model

Author: ricardoV94

pymc/dims/__init__.py

@@ -36,9 +36,15 @@ def __init__():
 
     # Make PyMC aware of xtensor functionality
     MeasurableOp.register(XRV)
-    lower_xtensor_query = optdb.query("+lower_xtensor")
-    logprob_rewrites_db.register("lower_xtensor", lower_xtensor_query, "basic", position=0.1)
-    initial_point_rewrites_db.register("lower_xtensor", lower_xtensor_query, "basic", position=0.1)
+    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
+    )
 
     # TODO: Better model of probability of bugs
     day_of_conception = datetime.date(2025, 6, 17)
@@ -64,4 +70,4 @@ def __init__():
 
 from pymc.dims import math
 from pymc.dims.distributions import *
-from pymc.dims.model import Data, Deterministic, Potential, with_dims
+from pymc.dims.model import Data, Deterministic, Potential

pymc/dims/distributions/core.py

@@ -14,15 +14,23 @@
 from collections.abc import Callable, Sequence
 from itertools import chain
 
+import numpy as np
+
+from pytensor.graph import node_rewriter
 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.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
+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
 
 
@@ -34,25 +42,97 @@ 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")
+
+    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)
+
+        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 = [dims[i] for i in inverse_transpose]
+            core_dims = dims_before_transpose[-ndim_supp:]
+        else:
+            core_dims = ()
+
+        return [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
+        return [MeasurableXTensorFromTensor(dims=node.op.dims, core_dims=None)(*node.inputs)]
+
+
+@_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: Callable | 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,
@@ -117,10 +197,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,
@@ -151,13 +243,16 @@ def dist(
         if dims_dict is None:
             extra_dims = None
         else:
-            parameter_implied_dims = set(
-                chain.from_iterable(param.type.dims for param in dist_params)
-            )
+            # Exclude dims that are implied by the parameters or core_dims
+            implied_dims = set(chain.from_iterable(param.type.dims for param in dist_params))
+            if core_dims is not None:
+                if isinstance(core_dims, str):
+                    implied_dims.add(core_dims)
+                else:
+                    implied_dims.update(core_dims)
+
             extra_dims = {
-                dim: length
-                for dim, length in dims_dict.items()
-                if dim not in parameter_implied_dims
+                dim: length for dim, length in dims_dict.items() if dim not in implied_dims
             }
         return cls.xrv_op(*dist_params, extra_dims=extra_dims, core_dims=core_dims, **kwargs)
 
@@ -177,10 +272,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/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
@@ -25,29 +24,6 @@
 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: DimsWithEllipsis = None, model: Model | None = None, **kwargs
 ) -> XTensorVariable:
@@ -79,8 +55,6 @@ 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
 
     value = registration_func(name, value, dims=dims, model=model)
 

pymc/dims/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/initial_point.py

@@ -20,7 +20,8 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Constant, Variable
+from pytensor.compile.ops import TypeCastingOp
+from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery, SequenceDB
 from pytensor.tensor.variable import TensorVariable
@@ -195,6 +196,14 @@ def inner(seed, *args, **kwargs):
     return make_seeded_function(func)
 
 
+class InitialPoint(TypeCastingOp):
+    def make_node(self, var):
+        return Apply(self, [var], [var.type()])
+
+
+initial_point_op = InitialPoint()
+
+
 def make_initial_point_expression(
     *,
     free_rvs: Sequence[TensorVariable],
@@ -235,6 +244,9 @@ def make_initial_point_expression(
 
     # Clone free_rvs so we don't modify the original graph
     initial_point_fgraph = FunctionGraph(outputs=free_rvs, clone=True)
+    # Wrap each rv in an initial_point Operation to avoid losing dependency between the RVs
+    replacements = tuple((rv, initial_point_op(rv)) for rv in initial_point_fgraph.outputs)
+    toposort_replace(initial_point_fgraph, replacements, reverse=True)
 
     # Apply any rewrites necessary to compute the initial points.
     initial_point_rewriter = initial_point_rewrites_db.query(initial_point_basic_query)
@@ -254,10 +266,10 @@ def make_initial_point_expression(
         if isinstance(strategy, str):
             if strategy == "support_point":
                 try:
-                    value = support_point(variable)
+                    value = support_point(variable.owner.inputs[0])
                 except NotImplementedError:
                     warnings.warn(
-                        f"Moment not defined for variable {variable} of type "
+                        f"support_point not defined for variable {variable} of type "
                         f"{variable.owner.op.__class__.__name__}, defaulting to "
                         f"a draw from the prior. This can lead to difficulties "
                         f"during tuning. You can manually define an initval or "