diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 1d9278dd34..6ada4ec369 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -142,6 +142,7 @@ jobs: - | tests/dims/distributions/test_core.py tests/dims/distributions/test_censored.py + tests/dims/distributions/test_custom.py tests/dims/distributions/test_scalar.py tests/dims/distributions/test_vector.py tests/dims/test_model.py diff --git a/pymc/dims/distributions/__init__.py b/pymc/dims/distributions/__init__.py index 6c49789089..fb87ee3a4c 100644 --- a/pymc/dims/distributions/__init__.py +++ b/pymc/dims/distributions/__init__.py @@ -12,5 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from pymc.dims.distributions.censored import Censored +from pymc.dims.distributions.custom import CustomDist from pymc.dims.distributions.scalar import * from pymc.dims.distributions.vector import * diff --git a/pymc/dims/distributions/censored.py b/pymc/dims/distributions/censored.py index 7fc1e0e02e..2fd0b6ff53 100644 --- a/pymc/dims/distributions/censored.py +++ b/pymc/dims/distributions/censored.py @@ -13,49 +13,89 @@ # limitations under the License. import numpy as np -from pytensor.xtensor.random.variable import shared_rng as xtensor_shared_rng +from pytensor.xtensor import as_xtensor, broadcast +from pytensor.xtensor.random import shared_rng +from pytensor.xtensor.random.type import xrandom_generator_type -from pymc.dims.distributions.core import DimDistribution, copy_docstring, expand_dist_dims +from pymc.dims.distributions.core import ( + DimDistribution, + DimSymbolicRandomVariable, + copy_docstring, + expand_dist_dims, +) from pymc.distributions.censored import Censored as RegularCensored +from pymc.distributions.censored import support_point_censored +from pymc.distributions.distribution import _support_point -@copy_docstring(RegularCensored) -class Censored(DimDistribution): - @classmethod - def dist(cls, dist, *, lower=None, upper=None, dim_lengths, **kwargs): - if lower is None: - lower = -np.inf - if upper is None: - upper = np.inf - return super().dist([dist, lower, upper], dim_lengths=dim_lengths, **kwargs) +class DimCensoredRV(DimSymbolicRandomVariable): + """Censored distribution on XTensorVariables. - @classmethod - def xrv_op(cls, dist, lower, upper, core_dims=None, extra_dims=None, rng=None, **kwargs): - if extra_dims is None: - extra_dims = {} + The logp is derived from the inner clip graph. The randomness belongs to + the censored dist input, so the RNG input is passed through unchanged. + """ - dist = cls._as_xtensor(dist) - lower = cls._as_xtensor(lower) - upper = cls._as_xtensor(upper) + inline_logprob = True + _print_name = ("Censored", "\\operatorname{Censored}") + + @classmethod + def rv_op( + cls, + dist, + lower, + upper, + *, + core_dims=None, + extra_dims=None, + rng=None, + # The next rng is always returned; the argument exists only while + # the pytensor XRV constructors transition to doing the same + return_next_rng: bool = True, + ): + assert return_next_rng, "return_next_rng=False is not supported" + dist = DimDistribution._as_xtensor(dist) + lower = DimDistribution._as_xtensor(lower) + upper = DimDistribution._as_xtensor(upper) # Any dimensions in extra_dims, or only present in lower, upper, # must propagate back to the dist as `extra_dims` bounds_sizes = lower.sizes | upper.sizes dist_dims_set = set(dist.dims) - extra_dist_dims = extra_dims | { + extra_dist_dims = (extra_dims or {}) | { dim: size for dim, size in bounds_sizes.items() if dim not in dist_dims_set } if extra_dist_dims: dist = expand_dist_dims(dist, extra_dist_dims) - # Probability is inferred from the clip operation - # TODO: Make this a SymbolicRandomVariable that can itself be resized - clipped = dist.clip(lower, upper) - if kwargs.get("return_next_rng"): - # TODO: Hack -- we have no rng of our own to thread forward. - # change_dist_size should grow a return_next_rng option so that - # Censored (and similar wrappers) can retrieve and forward the - # underlying dist's next_rng instead of returning a throwaway - # shared rng here. - return xtensor_shared_rng(seed=None), clipped - return clipped + # Censoring is achieved by clipping the dist between lower and upper. + # The randomness belongs to the dist input, so the RNG is passed through + dummy_rng = xrandom_generator_type("rng") + op = cls( + inputs=[dist, lower, upper, dummy_rng], + outputs=[dist.clip(lower, upper), dummy_rng], + ) + if rng is None: + rng = shared_rng(seed=None) + censored, next_rng = op(dist, lower, upper, rng) + return next_rng, censored + + +@_support_point.register(DimCensoredRV) +def dim_censored_support_point(op, rv, dist, lower, upper, rng): + # Align the inputs by name and reuse the regular CensoredRV implementation + dist, lower, upper = broadcast(dist, lower, upper) + sp = support_point_censored(op, rv.values, dist.values, lower.values, upper.values) + return as_xtensor(sp, dims=dist.type.dims).transpose(*rv.type.dims) + + +@copy_docstring(RegularCensored) +class Censored(DimDistribution): + xrv_op = DimCensoredRV.rv_op + + @classmethod + def dist(cls, dist, *, lower=None, upper=None, dim_lengths, **kwargs): + if lower is None: + lower = -np.inf + if upper is None: + upper = np.inf + return super().dist([dist, lower, upper], dim_lengths=dim_lengths, **kwargs) diff --git a/pymc/dims/distributions/core.py b/pymc/dims/distributions/core.py index 42a642b826..83a0a6490f 100644 --- a/pymc/dims/distributions/core.py +++ b/pymc/dims/distributions/core.py @@ -16,16 +16,20 @@ from typing import Any, cast import numpy as np -import pytensor.tensor as pt +from pytensor.compile.builders import OpFromGraph from pytensor.graph import node_rewriter from pytensor.graph.basic import Variable +from pytensor.graph.rewriting.basic import in2out from pytensor.tensor import TensorVariable from pytensor.tensor import expand_dims as pt_expand_dims from pytensor.tensor.elemwise import DimShuffle -from pytensor.tensor.random.op import RandomVariable +from pytensor.tensor.random.op import RandomVariable, RNGConsumerOp +from pytensor.tensor.random.type import RandomType +from pytensor.tensor.rewriting.ofg import inline_ofg_node from pytensor.xtensor import as_xtensor -from pytensor.xtensor.basic import XTensorFromTensor, xtensor_from_tensor +from pytensor.xtensor.basic import TensorFromXTensor, XTensorFromTensor, xtensor_from_tensor +from pytensor.xtensor.random import shared_rng from pytensor.xtensor.shape import Transpose from pytensor.xtensor.type import XTensorVariable from pytensor.xtensor.vectorization import XRV @@ -35,12 +39,96 @@ from pymc.distributions.distribution import _support_point, support_point from pymc.distributions.shape_utils import DimsWithEllipsis, convert_dims_with_ellipsis from pymc.logprob.abstract import MeasurableOp, _icdf, _logccdf, _logcdf, _logprob -from pymc.logprob.rewriting import measurable_ir_rewrites_db +from pymc.logprob.rewriting import logprob_rewrites_db, measurable_ir_rewrites_db from pymc.logprob.tensor import MeasurableDimShuffle from pymc.logprob.utils import filter_measurable_variables from pymc.util import UNSET +class DimSymbolicRandomVariable(MeasurableOp, RNGConsumerOp, OpFromGraph): + """Base class for dims distributions defined by an inner xtensor graph. + + The xtensor counterpart of :class:`~pymc.distributions.distribution.SymbolicRandomVariable`, + for dims distributions that need an Op to dispatch logp/logcdf/support_point + on, such as CustomDist and factory distributions like Censored or Truncated. + + It operates on XTensorVariables with named dims. Like a RandomVariable, the + Op has a single RNG input (the last input) and its outputs are the random + variable and the final RNG state. The inner graph must pipe the RNG through + its random operations, or pass it through unchanged if it has none. + + The Op is inlined away wherever the demarcation is not needed: at compile + time (``is_inline``) and during logprob inference when the logp is derived + automatically from the inner graph (``inline_logprob``). + """ + + inline_logprob = False + rv_op: Callable | None = None + """Constructor of the variable, with signature (*params, extra_dims, rng), + that returns (next_rng, rv) like the pytensor XRV constructors. + + Like in SymbolicRandomVariable, it is defined on the Op subclass: as a + classmethod, or as a staticmethod closure on dynamically created subclasses + when it requires build-time state (as in CustomDist). + """ + + def __init__(self, *args, extra_dims: Sequence[str] = (), **kwargs): + # The inputs must be (*params, *extra_dim_lengths, rng) + self.extra_dims = tuple(extra_dims) + kwargs.setdefault("inline", True) + kwargs.setdefault("on_unused_input", "ignore") + super().__init__(*args, **kwargs) + rng_inputs = [inp for inp in self.inner_inputs if isinstance(inp.type, RandomType)] + if len(rng_inputs) != 1 or self.inner_inputs[-1] is not rng_inputs[0]: + raise ValueError(f"{type(self).__name__} requires a single RNG as the last input") + if len(self.inner_outputs) != 2 or not isinstance(self.inner_outputs[-1].type, RandomType): + raise ValueError(f"{type(self).__name__} requires (rv, next_rng) as outputs") + + @property + def n_params(self) -> int: + return len(self.inner_inputs) - len(self.extra_dims) - 1 + + def update(self, node) -> dict[Variable, Variable]: + [rng_input] = [inp for inp in node.inputs if isinstance(inp.type, RandomType)] + return {rng_input: node.outputs[1]} + + def rebuild_with_extra_dims(self, node, extra_dims: dict[str, Any]) -> XTensorVariable: + """Recreate the variable of this node with additional extra batch dims. + + Used by `expand_dist_dims` when factory distributions need to add dims + to their components. A fresh RNG is used, so the new and old variables + are not correlated. + """ + if self.rv_op is None: + raise NotImplementedError( + f"{type(self).__name__} does not define rv_op and cannot be rebuilt" + ) + n_params = self.n_params + params = node.inputs[:n_params] + old_extra_dims = dict( + zip(self.extra_dims, node.inputs[n_params : n_params + len(self.extra_dims)]) + ) + _next_rng, rv = self.rv_op(*params, extra_dims={**extra_dims, **old_extra_dims}) + return rv + + +@node_rewriter([DimSymbolicRandomVariable]) +def inline_dim_symbolic_rv(fgraph, node): + """Expand a DimSymbolicRandomVariable when obtaining the logp graph if `inline_logprob` is True.""" + if not node.op.inline_logprob: + return None + return inline_ofg_node(node) + + +# Registered before pre-canonicalization, like inline_SymbolicRandomVariable +logprob_rewrites_db.register( + "inline_DimSymbolicRandomVariable", + in2out(inline_dim_symbolic_rv), + "basic", + position=-20, +) + + @_support_point.register(DimShuffle) def dimshuffle_support_point(ds_op, _, rv): # We implement support point for DimShuffle because @@ -176,6 +264,70 @@ def measurable_xtensor_from_tensor_icdf(op, value, rv, **kwargs): ) +class MeasurableTensorFromXTensor(MeasurableOp, TensorFromXTensor): + """Bridge that lets tensor-level measurable machinery see through the type boundary. + + Needed when a measurable xtensor variable (e.g. a DimSymbolicRandomVariable + that is not inlined) is consumed by lowered tensor operations, like the clip + of a Censored distribution. + """ + + # Some logprob implementations use the name of the base RV Op + name = None + + +@node_rewriter([TensorFromXTensor]) +def find_measurable_tensor_from_xtensor(fgraph, node) -> list[TensorVariable] | None: + if isinstance(node.op, MeasurableTensorFromXTensor): + return None + if not filter_measurable_variables(node.inputs): + return None + return [cast(TensorVariable, MeasurableTensorFromXTensor()(*node.inputs))] + + +def _lower_logp_term(term, x_dims: tuple[str, ...]) -> TensorVariable: + # Order the xtensor logp term by the positional layout of the value + if isinstance(term, XTensorVariable): + term = term.transpose(*(dim for dim in x_dims if dim in term.type.dims)).values + return term + + +@_logprob.register(MeasurableTensorFromXTensor) +def measurable_tensor_from_xtensor_logprob(op, values, x, **kwargs): + [value] = values + value_xt = xtensor_from_tensor(value, dims=x.type.dims) + logp = _logprob(x.owner.op, (value_xt,), *x.owner.inputs, **kwargs) + if isinstance(logp, tuple | list): + [logp] = logp + return _lower_logp_term(logp, x.type.dims) + + +@_logcdf.register(MeasurableTensorFromXTensor) +def measurable_tensor_from_xtensor_logcdf(op, value, x, **kwargs): + value_xt = xtensor_from_tensor(value, dims=x.type.dims) + logcdf = _logcdf(x.owner.op, value_xt, *x.owner.inputs, **kwargs) + return _lower_logp_term(logcdf, x.type.dims) + + +@_logccdf.register(MeasurableTensorFromXTensor) +def measurable_tensor_from_xtensor_logccdf(op, value, x, **kwargs): + value_xt = xtensor_from_tensor(value, dims=x.type.dims) + logccdf = _logccdf(x.owner.op, value_xt, *x.owner.inputs, **kwargs) + return _lower_logp_term(logccdf, x.type.dims) + + +@_icdf.register(MeasurableTensorFromXTensor) +def measurable_tensor_from_xtensor_icdf(op, value, x, **kwargs): + value_xt = xtensor_from_tensor(value, dims=x.type.dims) + icdf = _icdf(x.owner.op, value_xt, *x.owner.inputs, **kwargs) + return _lower_logp_term(icdf, x.type.dims) + + +measurable_ir_rewrites_db.register( + "measurable_tensor_from_xtensor", find_measurable_tensor_from_xtensor, "basic", "xtensor" +) + + def copy_docstring(regular_cls): # Copy docstring from regular distribution class to dims class def get_regular_docstring(dims_cls): @@ -324,7 +476,7 @@ def dist( dim: length for dim, length in dim_lengths.items() if dim not in implied_dims } if kwargs.get("rng") is None: - kwargs["rng"] = pt.random.shared_rng(seed=None) + kwargs["rng"] = shared_rng(seed=None) _, rv = cls.xrv_op( *dist_params, extra_dims=extra_dims, @@ -372,11 +524,13 @@ def expand_dist_dims(dist: XTensorVariable, extra_dims: dict[str, Any]) -> XTens new_dist_op = type(dist.owner.op)(**dist_props) _old_rng, *params_and_dim_lengths = dist.owner.inputs # We don't propagate the old RNG, because we don't want the new and old dists to be correlated - new_rng = pt.random.shared_rng(seed=None) + new_rng = shared_rng(seed=None) return new_dist_op(new_rng, *extra_dims.values(), *params_and_dim_lengths) case Transpose(): return expand_dist_dims(dist.owner.inputs[0], extra_dims=extra_dims).transpose( ..., *dist.dims ) + case DimSymbolicRandomVariable(): + return op.rebuild_with_extra_dims(dist.owner, extra_dims) case _: raise NotImplementedError(f"expand_dist_dims not implemented for {dist} with op {op}") diff --git a/pymc/dims/distributions/custom.py b/pymc/dims/distributions/custom.py new file mode 100644 index 0000000000..3e0c705db1 --- /dev/null +++ b/pymc/dims/distributions/custom.py @@ -0,0 +1,305 @@ +# Copyright 2026 - 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. +from collections.abc import Callable, Sequence +from typing import Any + +import pytensor.tensor as pt + +from pytensor.configdefaults import config +from pytensor.graph.basic import Apply +from pytensor.graph.fg import FunctionGraph +from pytensor.graph.op import Op +from pytensor.tensor.basic import infer_static_shape +from pytensor.xtensor import as_xtensor, broadcast +from pytensor.xtensor import zeros_like as x_zeros_like +from pytensor.xtensor.basic import xtensor_from_tensor +from pytensor.xtensor.random import shared_rng +from pytensor.xtensor.random.type import xrandom_generator_type +from pytensor.xtensor.type import XTensorVariable + +from pymc.dims.distributions.core import DimDistribution, DimSymbolicRandomVariable +from pymc.distributions.custom import default_not_implemented +from pymc.distributions.distribution import _support_point +from pymc.logprob.abstract import _logcdf, _logprob +from pymc.model.core import new_or_existing_block_model_access +from pymc.pytensorf import collect_default_updates + +BLOCK_MODEL_ACCESS_ERROR_MSG = ( + "Model variables cannot be created in the dist function. Use the `.dist` API" +) + +DIST_NOT_XTENSOR_ERROR_MSG = ( + "The `dist` function must return an XTensorVariable. " + "Use `pmd.Normal.dist(...)` or `xtensor_from_tensor(rv, dims=...)` " + "to ensure dims are attached to the output." +) + + +class NonRandomCustomDist(Op): + """Placeholder Op for a CustomDist defined without a dist function. + + Takes its shape as explicit scalar inputs and is only useful for its type + information; evaluating it raises. + """ + + __props__ = ("dtype",) + + def __init__(self, dtype: str): + self.dtype = config.floatX if dtype == "floatX" else dtype + + def make_node(self, *shape): + shape, static_shape = infer_static_shape(shape) + out = pt.tensor(dtype=self.dtype, shape=static_shape) + return Apply(self, list(shape), [out]) + + def infer_shape(self, fgraph, node, input_shapes): + return [node.inputs] + + def perform(self, node, inputs, output_storage): + raise NotImplementedError( + "Attempted to draw values from a CustomDist that was constructed without a " + "`dist` function. Please re-build your model and provide a callable to the " + "`dist` keyword argument to allow forward sampling." + ) + + +def _non_random_dist_maker(dtype: str) -> Callable: + """Create a placeholder dist for a CustomDist defined without a dist function. + + The returned function gives a placeholder variable the extra dims and the + batch dims of the params, so that everything but forward sampling works. + """ + + def non_random_dist(*args): + *xtensor_params, extra_dims = args + dim_lengths = dict(extra_dims) + for param in xtensor_params: + for dim, length in zip(param.type.dims, tuple(param.values.shape)): + dim_lengths.setdefault(dim, length) + placeholder = NonRandomCustomDist(dtype=dtype)(*dim_lengths.values()) + return xtensor_from_tensor(placeholder, dims=tuple(dim_lengths)) + + return non_random_dist + + +class DimCustomDistRV(DimSymbolicRandomVariable): + """Dims-native demarcation of a CustomDist random graph. + + The inputs are the params (with their own dims), the extra dim lengths, + and the piped RNG. User-provided logp/logcdf/support_point functions are + dispatched on subclasses of this Op, with xtensor semantics. + """ + + +class CustomDist(DimDistribution): + """Dims-aware CustomDist for pymc.dims. + + The ``dist`` function receives the parameters as ``XTensorVariable``, + followed by ``extra_dims``: a dict mapping the dims requested via ``dims`` + or ``observed`` that are not implied by any parameter to their lengths. + It is the counterpart of the ``size`` argument of :class:`~pymc.CustomDist`, + and the same argument used when implementing a new ``pymc.dims`` + distribution. The function must return an ``XTensorVariable`` random graph + with those dims, built from other ``pymc.dims`` distributions and/or + ``pytensor.xtensor`` operations:: + + import pytensor.xtensor.math as ptxm + + + def logitnormal_dist(mu, sigma, extra_dims): + return ptxm.sigmoid(pmd.Normal.dist(mu=mu, sigma=sigma, dim_lengths=extra_dims)) + + + with pm.Model(coords={"city": range(5)}): + pmd.CustomDist("x", mu, sigma, dist=logitnormal_dist, dims="city") + + When only ``dist`` is provided, the logp is derived automatically from the + graph, exactly as it is for the built-in ``pymc.dims`` distributions. + + User-provided ``logp``/``logcdf``/``support_point`` override the derived + ones. They are dispatched with xtensor semantics: the value, the random + variable and all the params are ``XTensorVariable`` with named dims, and + the functions may return an ``XTensorVariable`` (or a plain tensor). + + A distribution can also be defined without ``dist``, through ``logp`` + (and/or ``logcdf``, ``support_point``). The random graph is then a + placeholder with the extra dims and the batch dims of the parameters, + which raises if forward sampling is attempted:: + + with pm.Model(coords={"city": range(5)}): + pmd.CustomDist("y", mu, sigma, logp=normal_logp, dims="city") + """ + + @classmethod + def dist( + cls, + *dist_params, + dist: Callable | None = None, + logp: Callable | None = None, + logcdf: Callable | None = None, + support_point: Callable | None = None, + dtype: str = "floatX", + dim_lengths: dict | None = None, + core_dims: str | Sequence[str] | None = None, + **kwargs, + ): + kwargs.update( + dist=dist, + logp=logp, + logcdf=logcdf, + support_point=support_point, + dtype=dtype, + ) + return super().dist( + list(dist_params), + dim_lengths=dim_lengths, + core_dims=core_dims, + **kwargs, + ) + + @classmethod + def xrv_op( + cls, + *dist_params, + dist: Callable | None = None, + logp: Callable | None = None, + logcdf: Callable | None = None, + support_point: Callable | None = None, + dtype: str = "floatX", + class_name: str = "CustomDist", + core_dims: str | Sequence[str] | None = None, + extra_dims: dict[str, Any] | None = None, + rng=None, + # The next rng is always returned; the argument exists only while + # the pytensor XRV constructors transition to doing the same + return_next_rng: bool = True, + ): + assert return_next_rng, "return_next_rng=False is not supported" + # core_dims is not needed: all dims, core or batched, are known from + # the dist function output. + xtensor_params = [cls._as_xtensor(p) for p in dist_params] + extra_dims = extra_dims or {} + extra_dim_names = tuple(extra_dims) + + if dist is None: + if logp is None: + # Match the lazy failure of pm.CustomDist when the logp is requested + logp = default_not_implemented(class_name, "logp") + dist = _non_random_dist_maker(dtype) + + # Build the inner graph on dummy inputs. Like the size argument of + # pm.CustomDist, extra_dims (the dims requested via dims/observed that + # are not implied by any param) are passed as the last argument of the + # dist function, which is responsible for using them. + dummy_params = [param.type() for param in xtensor_params] + dummy_extra_lengths = [pt.scalar(f"{dim}_length", dtype="int64") for dim in extra_dim_names] + with new_or_existing_block_model_access(error_msg_on_access=BLOCK_MODEL_ACCESS_ERROR_MSG): + rv = dist(*dummy_params, dict(zip(extra_dim_names, dummy_extra_lengths))) + if not isinstance(rv, XTensorVariable): + raise TypeError(DIST_NOT_XTENSOR_ERROR_MSG) + if missing_dims := (set(extra_dim_names) - set(rv.type.dims)): + raise ValueError( + f"The `dist` function output is missing dims {sorted(missing_dims)}. " + "Dims that are not implied by the params must be added through the " + "`extra_dims` argument, as in " + "`pmd.Normal.dist(mu, sigma, dim_lengths=extra_dims)`." + ) + + # Pipe a single RNG through the inner random operations, chaining each + # one onto the state left by the previous, so that the Op has one RNG + # input and one final state output, like a RandomVariable + dummy_rng = xrandom_generator_type("rng") + updates = collect_default_updates( + inputs=[*dummy_params, *dummy_extra_lengths], outputs=(rv,) + ) + if updates: + fgraph = FunctionGraph(outputs=[rv, *updates.values()], clone=False) + # Chain in topological order of the consuming nodes, as a random + # operation may depend on the draws of another + node_order = {node: i for i, node in enumerate(fgraph.toposort())} + ordered_rngs = sorted( + updates, + key=lambda rng: min(node_order[client] for client, _ in fgraph.clients[rng]), + ) + chained_rngs = [dummy_rng, *(updates[rng] for rng in ordered_rngs)] + fgraph.replace_all(list(zip(ordered_rngs, chained_rngs[:-1])), import_missing=True) + next_rng = chained_rngs[-1] + else: + next_rng = dummy_rng + + def rv_op(*params, extra_dims, rng=None): + return cls.xrv_op( + *params, + dist=dist, + logp=logp, + logcdf=logcdf, + support_point=support_point, + class_name=class_name, + extra_dims=extra_dims, + rng=rng, + ) + + rv_type = type( + class_name, + (DimCustomDistRV,), + { + "inline_logprob": logp is None, + "_print_name": (class_name, f"\\operatorname{{{class_name}}}"), + "rv_op": staticmethod(rv_op), + }, + ) + + # ---- Dispatch the user overrides with xtensor semantics ---- + if logp is not None: + + @_logprob.register(rv_type) + def xcustom_dist_logp(op, values, *inputs, **kwargs): + [value] = values + return logp(value, *inputs[: op.n_params]) + + if logcdf is not None: + + @_logcdf.register(rv_type) + def xcustom_dist_logcdf(op, value, *inputs, **kwargs): + return logcdf(value, *inputs[: op.n_params]) + + @_support_point.register(rv_type) + def xcustom_dist_support_point(op, rv_out, *inputs): + params = inputs[: op.n_params] + if support_point is not None: + return support_point(rv_out, *params) + # The dims counterpart of the tensor template `pt.full(size, param)`: + # a zero broadcast to the params and extra dim lengths, so that the + # initial point never evaluates the random graph + zero = as_xtensor(pt.zeros((), dtype=rv_out.type.dtype)) + extra_lengths = inputs[op.n_params : op.n_params + len(op.extra_dims)] + if op.extra_dims: + zero = zero.expand_dims(dim=dict(zip(op.extra_dims, extra_lengths))) + out_dims = rv_out.type.dims + reduced_dims = {dim for p in params for dim in p.type.dims if dim not in out_dims} + sp, *_ = broadcast(zero, *params, exclude=sorted(reduced_dims)) + if set(sp.type.dims) != set(out_dims): + # Some output dims are only created inside the dist graph + return x_zeros_like(rv_out) + return sp.transpose(*out_dims) + + xop = rv_type( + inputs=[*dummy_params, *dummy_extra_lengths, dummy_rng], + outputs=[rv, next_rng], + extra_dims=extra_dim_names, + ) + if rng is None: + rng = shared_rng(seed=None) + out, out_next_rng = xop(*xtensor_params, *extra_dims.values(), rng) + return out_next_rng, out diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py index 258090fd82..b3f156436a 100644 --- a/pymc/distributions/distribution.py +++ b/pymc/distributions/distribution.py @@ -384,17 +384,23 @@ def __init__( super().__init__(*args, **kwargs) def make_node(self, *inputs): - # If we try to build the RV with a different size type (vector -> None or None -> vector) - # We need to rebuild the Op with new size type in the inner graph + # If we try to build the RV with input types that don't match the inner graph + # (such as a different size type, or params with other dimensionality) + # we need to rebuild the Op with an updated inner graph if self.extended_signature is not None: (rng_arg_idxs, size_arg_idx, param_idxs), _ = self.get_input_output_type_idxs( self.extended_signature ) if size_arg_idx is not None and len(rng_arg_idxs) == 1: - new_size_type = normalize_size_param(inputs[size_arg_idx]).type - if not self.input_types[size_arg_idx].is_super(new_size_type): - params = [inputs[idx] for idx in param_idxs] - size = inputs[size_arg_idx] + inputs = list(inputs) + size = inputs[size_arg_idx] = normalize_size_param(inputs[size_arg_idx]) + params = [as_tensor_variable(inputs[idx]) for idx in param_idxs] + for idx, param in zip(param_idxs, params): + inputs[idx] = param + if not self.input_types[size_arg_idx].is_super(size.type) or any( + not self.input_types[idx].is_super(param.type) + for idx, param in zip(param_idxs, params) + ): rng = inputs[rng_arg_idxs[0]] return self.rebuild_rv(*params, size=size, rng=rng).owner diff --git a/tests/dims/distributions/test_censored.py b/tests/dims/distributions/test_censored.py index 457d7827ed..ab58ce522f 100644 --- a/tests/dims/distributions/test_censored.py +++ b/tests/dims/distributions/test_censored.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import numpy as np +import pytensor.tensor as pt +import pytensor.xtensor.math as ptxm import pytest from pytensor.xtensor import as_xtensor @@ -20,7 +22,8 @@ import pymc.distributions as regular_distributions -from pymc.dims import Censored, Normal +from pymc import draw as pm_draw +from pymc.dims import Censored, CustomDist, Normal from pymc.model.core import Model from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph @@ -45,6 +48,13 @@ def test_censored_basic(lower, upper): assert_equivalent_random_graph(model, reference_model) assert_equivalent_logp_graph(model, reference_model) + # Dead branches of the support point switch constant-fold -inf + inf + with np.errstate(invalid="ignore"): + np.testing.assert_allclose( + model.initial_point()["y"], + reference_model.initial_point()["y"], + ) + def test_censored_dims(): """Test that both censored (and the underlying dist) have all the implied and explicit dims.""" @@ -88,3 +98,54 @@ def test_censored_dims(): c3_dist = c3.owner.inputs[0] assert isinstance(c3_dist.owner.op, XRV) assert c3_dist.dims == ("d", "c", "a", "b") + + +def test_censored_custom_dist(): + """Factory distributions compose with an overridden CustomDist through its Op.""" + from scipy import stats + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return -0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi)) + + def normal_logcdf(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.log(pt.erf((value - mu) / (sigma * pt.sqrt(2.0))) + 1.0) - pt.log(2.0) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + base = CustomDist.dist(0.0, 1.0, dist=normal_dist, logp=normal_logp, logcdf=normal_logcdf) + # The extra dim is added by rebuilding the CustomDist through its Op + Censored("y", base, lower=-1.0, upper=1.0, dims="city") + + draws = pm_draw(model["y"], draws=100, random_seed=1) + assert draws.shape == (100, 5) + assert ((draws >= -1) & (draws <= 1)).all() + + test_value = np.array([-1.0, -0.5, 0.0, 0.5, 1.0]) + logp_value = model.compile_logp()({"y": test_value}) + ref = stats.norm(0, 1) + expected = np.log(ref.cdf(-1)) + ref.logpdf([-0.5, 0.0, 0.5]).sum() + np.log(ref.sf(1)) + np.testing.assert_allclose(logp_value, expected) + + +def test_censored_custom_dist_derived(): + """Censored resizes a compound CustomDist without overrides through its Op.""" + + def logitnormal_dist(mu, sigma, extra_dims): + return ptxm.sigmoid(Normal.dist(mu, sigma, dim_lengths=extra_dims)) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + base = CustomDist.dist(0.0, 1.0, dist=logitnormal_dist) + Censored("y", base, lower=0.2, upper=0.8, dims="city") + + draws = pm_draw(model["y"], draws=100, random_seed=1) + assert draws.shape == (100, 5) + assert ((draws >= 0.2) & (draws <= 0.8)).all() + + # The support point comes from the bounds, like in regular Censored + np.testing.assert_allclose(model.initial_point()["y"], np.full(5, 0.5)) diff --git a/tests/dims/distributions/test_custom.py b/tests/dims/distributions/test_custom.py new file mode 100644 index 0000000000..3f402ab545 --- /dev/null +++ b/tests/dims/distributions/test_custom.py @@ -0,0 +1,427 @@ +# Copyright 2026 - 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 numpy as np +import pytensor.tensor as pt +import pytensor.xtensor.math as ptxm +import pytest + +from pytensor.xtensor import as_xtensor, full_like + +import pymc.distributions as regular_distributions + +from pymc import draw as pm_draw +from pymc.dims import CustomDist, Normal, Poisson +from pymc.model.core import Model +from tests.dims.utils import assert_equivalent_logp_graph, assert_equivalent_random_graph + +pytestmark = pytest.mark.filterwarnings( + "error", + r"ignore:^Numba will use object mode to run.*perform method\.:UserWarning", +) + + +class TestCustomDistSymbolic: + """Tests for the symbolic (dist=) path of pmd.CustomDist.""" + + def test_compound_non_xrv_output(self): + """Compound dist with non-XRV output gets extra dims via the extra_dims argument.""" + + def logitnormal_dist(mu, sigma, extra_dims): + return ptxm.sigmoid(Normal.dist(mu=mu, sigma=sigma, dim_lengths=extra_dims)) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + x = CustomDist("x", 0, 1, dist=logitnormal_dist, dims="city") + + assert set(x.dims) == {"city"} + + draws = pm_draw(model["x"], draws=5) + assert draws.shape == (5, 5) + assert (draws > 0).all() and (draws < 1).all() + + def test_basic(self): + """Symbolic path: dist function wrapping Normal.dist, compared against regular Normal.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist("x", 0, 1, dist=normal_dist, dims="city") + + with Model(coords=coords) as reference_model: + regular_distributions.Normal("x", 0, 1, dims="city") + + assert_equivalent_random_graph(model, reference_model) + assert_equivalent_logp_graph(model, reference_model) + + def test_param_dims_propagate(self): + """Params with dims propagate to the output.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + coords = {"city": range(5)} + mu = as_xtensor(np.array([0, 1, 2, 3, 4]), dims=("city",)) + sigma = as_xtensor(np.array([1, 2, 3, 4, 5]), dims=("city",)) + + with Model(coords=coords) as model: + x = CustomDist("x", mu, sigma, dist=normal_dist) + + assert set(x.dims) == {"city"} + assert x.type.shape == (5,) + + def test_dist_must_use_extra_dims(self): + """A dist function that ignores extra_dims fails when extra dims are requested.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma) + + coords = {"city": range(5)} + with Model(coords=coords): + with pytest.raises(ValueError, match="output is missing dims"): + CustomDist("x", 0, 1, dist=normal_dist, dims="city") + + +class TestCustomDistArbitrary: + """Tests for the arbitrarily-defined (logp=) path of pmd.CustomDist.""" + + def test_logp_only_cannot_sample(self): + """Forward sampling a CustomDist defined only through logp raises.""" + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist("x", 0, 1, logp=normal_logp, dims="city") + + with pytest.raises(NotImplementedError, match="CustomDist that was constructed without"): + pm_draw(model["x"]) + + def test_logp_basic(self): + """Arbitrary path with logp function and dims on output.""" + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + rng = np.random.default_rng(42) + observed = as_xtensor(rng.normal(0, 1, size=5), dims=("city",)) + + with Model(coords=coords) as model: + CustomDist( + "x", + 0, + 1, + logp=normal_logp, + observed=observed, + dims="city", + ) + + # Test that logp evaluates without error and returns finite values + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + def test_hybrid_dist_logp(self): + """Hybrid path: dist for sampling + logp override.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist( + "x", + 0, + 1, + dist=normal_dist, + logp=normal_logp, + dims="city", + ) + + # Verify sampling works (via draw) + draws = pm_draw(model["x"], draws=3) + assert draws.shape == (3, 5) + + # Verify logp evaluates + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + def test_hybrid_derived_params(self): + """Hybrid path: dist derives params.""" + + def poisson_dist(a, b, c, extra_dims): + lam = a + b + c + return Poisson.dist(mu=lam, dim_lengths=extra_dims) + + def poisson_logp(value, a, b, c): + value = value.values + lam = (a + b + c).values + return pt.sum(value * pt.log(lam) - lam - pt.gammaln(value + 1)) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist( + "x", + 0.5, + 0.3, + 0.2, + dist=poisson_dist, + logp=poisson_logp, + dims="city", + ) + + # pm.draw — evaluates the full dist graph (compound inference) + draws = pm_draw(model["x"], draws=3) + assert draws.shape == (3, 5) + assert draws.dtype.kind == "i" + + # logp evaluates + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + def test_hybrid_logp_override(self): + """Hybrid path: verify user logp overrides auto-derived logp.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def scaled_logp(value, mu, sigma): + """Custom logp that multiplies normal logp by 2.""" + value, mu, sigma = value.values, mu.values, sigma.values + normal_logp = pt.sum( + -0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi)) + ) + return 2.0 * normal_logp + + coords = {"city": range(5)} + with Model(coords=coords) as model_hybrid: + CustomDist( + "x", + 0, + 1, + dist=normal_dist, + logp=scaled_logp, + dims="city", + ) + + # Auto-derived logp (no logp override) + with Model(coords=coords) as model_auto: + CustomDist( + "x", + 0, + 1, + dist=normal_dist, + dims="city", + ) + + ip = model_hybrid.initial_point() + hybrid_logp = model_hybrid.compile_logp()(ip) + auto_logp = model_auto.compile_logp()(ip) + # Hybrid logp should be 2x auto logp + np.testing.assert_allclose(hybrid_logp, 2.0 * auto_logp) + + def test_hybrid_basic_dims(self): + """Hybrid path with dims on params.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + mu = as_xtensor(np.array([0.0, 0.5, 1.0, 1.5, 2.0]), dims=("city",)) + sigma = as_xtensor(np.array([1.0, 1.1, 1.2, 1.3, 1.4]), dims=("city",)) + + with Model(coords=coords) as model: + x = CustomDist("x", mu, sigma, dist=normal_dist, logp=normal_logp) + + assert set(x.dims) == {"city"} + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + draws = pm_draw(model["x"], draws=3) + assert draws.shape == (3, 5) + + def test_logcdf(self): + """Arbitrary path with logcdf function.""" + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + def normal_logcdf(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum( + pt.log(pt.erf((value - mu) / (sigma * pt.sqrt(2.0))) + 1.0) + - pt.log(pt.constant(2.0)) + ) + + coords = {"city": range(5)} + rng = np.random.default_rng(42) + observed = as_xtensor(rng.normal(0, 1, size=5), dims=("city",)) + + with Model(coords=coords) as model: + CustomDist( + "x", + 0, + 1, + logp=normal_logp, + logcdf=normal_logcdf, + observed=observed, + dims="city", + ) + + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + def test_mu_as_model_var(self): + """Arbitrary path with mu as a model variable (no dims on mu).""" + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + rng = np.random.default_rng(42) + observed = as_xtensor(rng.normal(0, 1, size=5), dims=("city",)) + + with Model(coords=coords) as model: + mu = Normal("mu", 0, 1) + CustomDist( + "x", + mu, + 1, + logp=normal_logp, + observed=observed, + dims="city", + ) + + ip = model.initial_point() + logp_val = model.compile_logp()(ip) + assert np.isfinite(logp_val) + + def test_logp_inferred_with_other_overrides(self): + """The logp is still derived from the dist graph when only other methods are overridden.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def custom_support_point(rv, mu, sigma): + return full_like(rv, mu) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist( + "x", + 1.0, + 2.0, + dist=normal_dist, + support_point=custom_support_point, + dims="city", + ) + + with Model(coords=coords) as reference_model: + regular_distributions.Normal("x", 1.0, 2.0, dims="city") + + ip = model.initial_point() + np.testing.assert_allclose(ip["x"], np.ones(5)) + np.testing.assert_allclose( + model.compile_logp()(ip), + reference_model.compile_logp()(ip), + ) + + def test_default_support_point(self): + """Default support point of a CustomDist without a dist function.""" + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist( + "x", + 0, + 1, + logp=normal_logp, + dims="city", + ) + + # The initial point does not require evaluating the random graph, + # which would raise for a CustomDist without a dist function + ip = model.initial_point() + np.testing.assert_allclose(ip["x"], np.zeros(5)) + + def test_default_support_point_reduced_param_dim(self): + """Default support point drops param dims that the dist graph reduced away.""" + + def averaged_normal_dist(mu, extra_dims): + return Normal.dist(mu.mean("city"), dim_lengths=extra_dims) + + def normal_logp(value, mu): + value, mu_mean = value.values, mu.mean("city").values + return pt.sum(-0.5 * (value - mu_mean) ** 2 - pt.log(pt.sqrt(2 * np.pi))) + + coords = {"city": range(5), "obs": range(3)} + mu = as_xtensor(np.arange(5, dtype="float64"), dims=("city",)) + + with Model(coords=coords) as model: + CustomDist("x", mu, dist=averaged_normal_dist, logp=normal_logp, dims="obs") + + ip = model.initial_point() + np.testing.assert_allclose(ip["x"], np.zeros(3)) + + def test_hybrid_support_point(self): + """Hybrid path with custom support_point.""" + + def normal_dist(mu, sigma, extra_dims): + return Normal.dist(mu, sigma, dim_lengths=extra_dims) + + def normal_logp(value, mu, sigma): + value, mu, sigma = value.values, mu.values, sigma.values + return pt.sum(-0.5 * ((value - mu) / sigma) ** 2 - pt.log(sigma * pt.sqrt(2 * np.pi))) + + def custom_support_point(rv, mu, sigma): + return full_like(rv, mu) + + coords = {"city": range(5)} + with Model(coords=coords) as model: + CustomDist( + "x", + 0, + 1, + dist=normal_dist, + logp=normal_logp, + support_point=custom_support_point, + dims="city", + ) + + # Like all dims distributions, support_point is reachable after lowering + ip = model.initial_point() + np.testing.assert_allclose(ip["x"], np.zeros(5)) diff --git a/tests/dims/utils.py b/tests/dims/utils.py index d873ef54e9..fb8dc3ae92 100644 --- a/tests/dims/utils.py +++ b/tests/dims/utils.py @@ -11,25 +11,52 @@ # 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.tensor as pt + from pytensor.graph import rewrite_graph +from pytensor.graph.traversal import ancestors from pytensor.graph.replace import graph_replace from pytensor.printing import debugprint from pytensor.xtensor import as_xtensor +from pytensor.xtensor.random.type import XRNGToRNG from pymc import Model from pymc.testing import equal_computations_up_to_root +def _normalize_shared_rngs(outputs): + """Replace root XRNGToRNG(shared_xrng) casts by plain tensor shared RNGs. + + Lowering leaves a cast between the shared xtensor RNG and the tensor RV, + which the reference tensor-based models don't have. + """ + replacements = { + var: pt.random.shared_rng(seed=None) + for var in ancestors(outputs) + if ( + var.owner is not None + and isinstance(var.owner.op, XRNGToRNG) + and var.owner.inputs[0].owner is None + ) + } + if not replacements: + return outputs + return graph_replace(outputs, replacements) + + def assert_equivalent_random_graph(model: Model, reference_model: Model) -> bool: """Check if the random graph of a model with xtensor variables is equivalent.""" - lowered_model = rewrite_graph( - [var.values for var in model.basic_RVs + model.deterministics + model.potentials], - include=( - "lower_xtensor", - "inline_ofg_expansion_xtensor", - "canonicalize", - "local_remove_all_assert", - ), + lowered_model = _normalize_shared_rngs( + rewrite_graph( + [var.values for var in model.basic_RVs + model.deterministics + model.potentials], + include=( + "inline_ofg_expansion", + "lower_xtensor", + "inline_ofg_expansion_xtensor", + "canonicalize", + "local_remove_all_assert", + ), + ) ) reference_lowered_model = rewrite_graph( reference_model.basic_RVs + reference_model.deterministics + reference_model.potentials,