Repository: pymc-devs/pymc
Pull request: Refactor DensityDist into v4
Closes #4831

This PR ports `DensityDist` into v4. It creates a `DensityDistRV` class and `DensityDist` class in a similar way as is done by the `Simulator` class. When a `DensityDist` is created, it registers the `logp`, `logcdf` and `get_moment` functions for the associated Op type. The `DensityDistRV` is used to overload the `rng_fn` from the base class to be able to run the user supplied `random` callable with the appropriate signature.

Depending on what your PR does, here are a few things you might want to address in the description:
+ [X] what are the (breaking) changes that this PR makes?

The signature of `DensityDist` is completely different than in v3:

- It now takes as positional arguments, the distribution parameters. In v3, these were passed in through a dictionary that was supplied as `observed`. The associated `MultiObservedRV` is no longer used or needed and can be removed in a following PR. 
- `logp` is now an optional keyword only argument instead of the first positional argument. If users try to pass it as the first positional argument, they will get an exception that tells them that they are using the old v3 API and should change.
- The signature of `logp` is now different, it takes in the distribution parameters as positional arguments. Something similar happens for `random`.
- All the distribution methods: `logp`, `random`, `logcdf` and `get_moment` are optional. If they are not passed during construction, the RV will be built, but a `NotImplementedError` will be raised when you try to do something like `pm.sample`, `pm.sample_prior_predictive` or `get_moment`.

+ [X] important background, or details about the implementation
+ [X] are the changes—especially new features—covered by tests and docstrings?

The docstrings were updated and new tests have been added for the different optional methods (except `logcdf`) which behaves in the same way as `logp`. If necessary, it's possible to add tests for `logcdf` too.

+ [X] [linting/style checks have been run](https://github.com/pymc-devs/pymc3/wiki/PyMC3-Python-Code-Style)
+ [x] [consider adding/updating relevant example notebooks](https://github.com/pymc-devs/pymc-examples)
+ [x] right before it's ready to merge, mention the PR in the RELEASE-NOTES.md


Problem statement:
Refactor or deprecate `DensityDist`
We need to decide whether we want to deprecate `DensityDist` and if so, provide at least documentation on how to add simple logp terms to a model or how to quickly implement a custom Distribution.

If we want to refactor it, it is probably good time to get rid of the whole dictionary as observed thing and force the distribution parameters to go into the signature as normal arguments in other distributions

```python
# V3 allows this
def my_logp(value, param1, param2):
    return value * param1 - param2

param1 = pm.Normal('param1', 0, 1)
param2 = pm.Normal('param2', 0, 1)
y = pm.DensityDist('y', my_logp, observed=dict(value=data, param1=param1, param2=param2))
```

```python
# V4 should work only like this (if not deprecated)
y = pm.DensityDist('y', my_logp, param1, param2, obeserved=data)

Hints from issue discussion:
CC @lucianopaz, if you want to assign this to yourself. Seems like we are pretty close to having a working solution
looks related to https://github.com/pymc-devs/pymc3/issues/4534, and I also wanted to leave a link to https://discourse.pymc.io/t/using-a-random-variable-as-observed/7184/5 so we are aware of all the (in my opinion strange and redundant) things DensityDist allows, expecially when taking a dictionary as `observed` (which I also think is the only Distribution that does). 

As I have commented a few times and started changing in https://github.com/pymc-devs/pymc3/pull/4433, I think it's a bad idea to allow `FreeRVs` to be values in the dictionary passed as observed. If this is a feature that is really important to DensityDist, I would suggest doing this `givens` extra argument or following @ricardoV94 suggestion.

Commit message: Refactor DensityDist into v4

Patch to review:
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 97261f82d8..742f5d9462 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -10,6 +10,9 @@
 - `pm.sample` now returns results as `InferenceData` instead of `MultiTrace` by default (see [#4744](https://github.com/pymc-devs/pymc/pull/4744)).
 - `pm.sample_prior_predictive` no longer returns transformed variable values by default. Pass them by name in `var_names` if you want to obtain these draws (see [4769](https://github.com/pymc-devs/pymc/pull/4769)).
 - ⚠ `pm.Bound` interface no longer accepts a callable class as argument, instead it requires an instantiated distribution (created via the `.dist()` API) to be passed as an argument. In addition, Bound no longer returns a class instance but works as a normal PyMC distribution. Finally, it is no longer possible to do predictive random sampling from Bounded variables. Please, consult the new documentation for details on how to use Bounded variables (see [4815](https://github.com/pymc-devs/pymc/pull/4815)).
+- `pm.DensityDist` no longer accepts the `logp` as its first position argument. It is now an optional keyword argument. If you pass a callable as the first positional argument, a `TypeError` will be raised (see [5026](https://github.com/pymc-devs/pymc3/pull/5026)).
+- `pm.DensityDist` now accepts distribution parameters as positional arguments. Passing them as a dictionary in the `observed` keyword argument is no longer supported and will raise an error (see [5026](https://github.com/pymc-devs/pymc3/pull/5026)).
+- The signature of the `logp` and `random` functions that can be passed into a `pm.DensityDist` has been changed (see [5026](https://github.com/pymc-devs/pymc3/pull/5026)).
 - ...
 
 ### New Features
@@ -25,6 +28,8 @@
 - The `Polya-Gamma` distribution has been added (see [#4531](https://github.com/pymc-devs/pymc/pull/4531)). To make use of this distribution, the [`polyagamma>=1.3.1`](https://pypi.org/project/polyagamma/) library must be installed and available in the user's environment.
 - A small change to the mass matrix tuning methods jitter+adapt_diag (the default) and adapt_diag improves performance early on during tuning for some models. [#5004](https://github.com/pymc-devs/pymc/pull/5004)
 - New experimental mass matrix tuning method jitter+adapt_diag_grad. [#5004](https://github.com/pymc-devs/pymc/pull/5004)
+- `pm.DensityDist` can now accept an optional `logcdf` keyword argument to pass in a function to compute the cummulative density function of the distribution (see [5026](https://github.com/pymc-devs/pymc3/pull/5026)).
+- `pm.DensityDist` can now accept an optional `get_moment` keyword argument to pass in a function to compute the moment of the distribution (see [5026](https://github.com/pymc-devs/pymc3/pull/5026)).
 - ...
 
 ### Maintenance
diff --git a/docs/source/Probability_Distributions.rst b/docs/source/Probability_Distributions.rst
index 63d339e499..7b8999e487 100644
--- a/docs/source/Probability_Distributions.rst
+++ b/docs/source/Probability_Distributions.rst
@@ -58,16 +58,16 @@ An exponential survival function, where :math:`c=0` denotes failure (or non-surv
     f(c, t) = \left\{ \begin{array}{l} \exp(-\lambda t), \text{if c=1} \\
                \lambda \exp(-\lambda t), \text{if c=0}  \end{array} \right.
 
-Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as an argument to the ``DensityDist`` function, which creates an instance of a PyMC distribution with the custom function as its log-probability.
+Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as a keyword argument to the ``DensityDist`` function, which creates an instance of a PyMC3 distribution with the custom function as its log-probability.
 
 For the exponential survival function, this is:
 
 ::
 
-    def logp(failure, value):
-        return (failure * log(λ) - λ * value).sum()
+    def logp(value, t, λ):
+        return (value * log(λ) - λ * t).sum()
 
-    exp_surv = pm.DensityDist('exp_surv', logp, observed={'failure':failure, 'value':t})
+    exp_surv = pm.DensityDist('exp_surv', t, λ, logp=logp, observed=failure)
 
 Similarly, if a random number generator is required, a function returning random numbers corresponding to the probability distribution can be passed as the ``random`` argument.
 
diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
index 2ffa092cd6..3b9b91593e 100644
--- a/pymc/distributions/distribution.py
+++ b/pymc/distributions/distribution.py
@@ -13,17 +13,18 @@
 #   limitations under the License.
 import contextvars
 import functools
-import multiprocessing
 import sys
 import types
 import warnings
 
 from abc import ABCMeta
 from functools import singledispatch
-from typing import Optional
+from typing import Callable, Optional, Sequence
 
 import aesara
+import numpy as np
 
+from aesara.tensor.basic import as_tensor_variable
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
 from aesara.tensor.var import TensorVariable
@@ -41,12 +42,14 @@
     maybe_resize,
     resize_from_dims,
     resize_from_observed,
+    to_tuple,
 )
 from pymc.printing import str_for_dist
 from pymc.util import UNSET
 from pymc.vartypes import string_types
 
 __all__ = [
+    "DensityDistRV",
     "DensityDist",
     "Distribution",
     "Continuous",
@@ -387,96 +390,241 @@ class NoDistribution(Distribution):
     """
 
 
-class DensityDist(Distribution):
-    """Distribution based on a given log density function.
+class DensityDistRV(RandomVariable):
+    """
+    Base class for DensityDistRV
+
+    This should be subclassed when defining custom DensityDist objects.
+    """
+
+    name = "DensityDistRV"
+    _print_name = ("DensityDist", "\\operatorname{DensityDist}")
+
+    @classmethod
+    def rng_fn(cls, rng, *args):
+        args = list(args)
+        size = args.pop(-1)
+        return cls._random_fn(*args, rng=rng, size=size)
 
-    A distribution with the passed log density function is created.
-    Requires a custom random function passed as kwarg `random` to
-    enable prior or posterior predictive sampling.
 
+class DensityDist(NoDistribution):
+    """A distribution that can be used to wrap black-box log density functions.
+
+    Creates a Distribution and registers the supplied log density function to be used
+    for inference. It is also possible to supply a `random` method in order to be able
+    to sample from the prior or posterior predictive distributions.
     """
 
-    def __init__(
-        self,
-        logp,
-        shape=(),
-        dtype=None,
-        initval=0,
-        random=None,
-        wrap_random_with_dist_shape=True,
-        check_shape_in_random=True,
-        *args,
+    def __new__(
+        cls,
+        name: str,
+        *dist_params,
+        logp: Optional[Callable] = None,
+        logcdf: Optional[Callable] = None,
+        random: Optional[Callable] = None,
+        get_moment: Optional[Callable] = None,
+        ndim_supp: int = 0,
+        ndims_params: Optional[Sequence[int]] = None,
+        dtype: str = "floatX",
         **kwargs,
     ):
         """
         Parameters
         ----------
-
-        logp: callable
-            A callable that has the following signature ``logp(value)`` and
-            returns an Aesara tensor that represents the distribution's log
-            probability density.
-        shape: tuple (Optional): defaults to `()`
-            The shape of the distribution. The default value indicates a scalar.
-            If the distribution is *not* scalar-valued, the programmer should pass
-            a value here.
-        dtype: None, str (Optional)
-            The dtype of the distribution.
-        initval: number or array (Optional)
-            The ``initval`` of the RV's tensor that follow the ``DensityDist``
-            distribution.
-        args, kwargs: (Optional)
-            These are passed to the parent class' ``__init__``.
+        name : str
+        dist_params : Tuple
+            A sequence of the distribution's parameter. These will be converted into
+            aesara tensors internally. These parameters could be other ``RandomVariable``
+            instances.
+        logp : Optional[Callable]
+            A callable that calculates the log density of some given observed ``value``
+            conditioned on certain distribution parameter values. It must have the
+            following signature: ``logp(value, *dist_params)``, where ``value`` is
+            an Aesara tensor that represents the observed value, and ``dist_params``
+            are the tensors that hold the values of the distribution parameters.
+            This function must return an Aesara tensor. If ``None``, a ``NotImplemented``
+            error will be raised when trying to compute the distribution's logp.
+        logcdf : Optional[Callable]
+            A callable that calculates the log cummulative probability of some given observed
+            ``value`` conditioned on certain distribution parameter values. It must have the
+            following signature: ``logcdf(value, *dist_params)``, where ``value`` is
+            an Aesara tensor that represents the observed value, and ``dist_params``
+            are the tensors that hold the values of the distribution parameters.
+            This function must return an Aesara tensor. If ``None``, a ``NotImplemented``
+            error will be raised when trying to compute the distribution's logcdf.
+        random : Optional[Callable]
+            A callable that can be used to generate random draws from the distribution.
+            It must have the following signature: ``random(*dist_params, rng=None, size=None)``.
+            The distribution parameters are passed as positional arguments in the
+            same order as they are supplied when the ``DensityDist`` is constructed.
+            The keyword arguments are ``rnd``, which will provide the random variable's
+            associated :py:class:`~numpy.random.Generator`, and ``size``, that will represent
+            the desired size of the random draw. If ``None``, a ``NotImplemented``
+            error will be raised when trying to draw random samples from the distribution's
+            prior or posterior predictive.
+        get_moment : Optional[Callable]
+            A callable that can be used to compute the moments of the distribution.
+            It must have the following signature: ``get_moment(rv, size, *rv_inputs)``.
+            The distribution's :class:`~aesara.tensor.random.op.RandomVariable` is passed
+            as the first argument ``rv``. ``size`` is the random variable's size implied
+            by the ``dims``, ``size`` and parameters supplied to the distribution. Finally,
+            ``rv_inputs`` is the sequence of the distribution parameters, in the same order
+            as they were supplied when the DensityDist was created. If ``None``, a
+            ``NotImplemented`` error will be raised when trying to draw random samples from
+            the distribution's prior or posterior predictive.
+        ndim_supp : int
+            The number of dimensions in the support of the distribution. Defaults to assuming
+            a scalar distribution, i.e. ``ndim_supp = 0``.
+        ndims_params : Optional[Sequence[int]]
+            The list of number of dimensions in the support of each of the distribution's
+            parameters. If ``None``, it is assumed that all parameters are scalars, hence
+            the number of dimensions of their support will be 0.
+        dtype : str
+            The dtype of the distribution. All draws and observations passed into the distribution
+            will be casted onto this dtype.
+        kwargs :
+            Extra keyword arguments are passed to the parent's class ``__new__`` method.
 
         Examples
         --------
             .. code-block:: python
 
+                def logp(value, mu):
+                    return -(value - mu)**2
+
                 with pm.Model():
                     mu = pm.Normal('mu',0,1)
-                    normal_dist = pm.Normal.dist(mu, 1)
                     pm.DensityDist(
                         'density_dist',
-                        normal_dist.logp,
+                        mu,
+                        logp=logp,
                         observed=np.random.randn(100),
                     )
                     idata = pm.sample(100)
 
             .. code-block:: python
 
+                def logp(value, mu):
+                    return -(value - mu)**2
+
+                def random(mu, rng=None, size=None):
+                    return rng.normal(loc=mu, scale=1, size=size)
+
                 with pm.Model():
                     mu = pm.Normal('mu', 0 , 1)
-                    normal_dist = pm.Normal.dist(mu, 1, shape=3)
                     dens = pm.DensityDist(
                         'density_dist',
-                        normal_dist.logp,
+                        mu,
+                        logp=logp,
+                        random=random,
                         observed=np.random.randn(100, 3),
-                        shape=3,
+                        size=(100, 3),
                     )
                     prior = pm.sample_prior_predictive(10)['density_dist']
                 assert prior.shape == (10, 100, 3)
 
         """
-        if dtype is None:
+
+        if dist_params is None:
+            dist_params = []
+        elif len(dist_params) > 0 and callable(dist_params[0]):
+            raise TypeError(
+                "The DensityDist API has changed, you are using the old API "
+                "where logp was the first positional argument. In the current API, "
+                "the logp is a keyword argument, amongst other changes. Please refer "
+                "to the API documentation for more information on how to use the "
+                "new DensityDist API."
+            )
+        dist_params = [as_tensor_variable(param) for param in dist_params]
+
+        # Assume scalar ndims_params
+        if ndims_params is None:
+            ndims_params = [0] * len(dist_params)
+
+        if logp is None:
+            logp = default_not_implemented(name, "logp")
+
+        if logcdf is None:
+            logcdf = default_not_implemented(name, "logcdf")
+
+        if random is None:
+            random = default_not_implemented(name, "random")
+
+        if get_moment is None:
+            get_moment = default_not_implemented(name, "get_moment")
+
+        rv_op = type(
+            f"DensityDist_{name}",
+            (DensityDistRV,),
+            dict(
+                name=f"DensityDist_{name}",
+                inplace=False,
+                ndim_supp=ndim_supp,
+                ndims_params=ndims_params,
+                dtype=dtype,
+                # Specifc to DensityDist
+                _random_fn=random,
+            ),
+        )()
+
+        # Register custom logp
+        rv_type = type(rv_op)
+
+        @_logp.register(rv_type)
+        def density_dist_logp(op, rv, rvs_to_values, *dist_params, **kwargs):
+            value_var = rvs_to_values.get(rv, rv)
+            return logp(
+                value_var,
+                *dist_params,
+            )
+
+        @_logcdf.register(rv_type)
+        def density_dist_logcdf(op, var, rvs_to_values, *dist_params, **kwargs):
+            value_var = rvs_to_values.get(var, var)
+            return logcdf(value_var, *dist_params, **kwargs)
+
+        @_get_moment.register(rv_type)
+        def density_dist_get_moment(op, rv, size, *rv_inputs):
+            return get_moment(rv, size, *rv_inputs)
+
+        cls.rv_op = rv_op
+        return super().__new__(cls, name, *dist_params, **kwargs)
+
+    @classmethod
+    def dist(cls, *args, **kwargs):
+        output = super().dist(args, **kwargs)
+        if cls.rv_op.dtype == "floatX":
             dtype = aesara.config.floatX
-        super().__init__(shape, dtype, initval, *args, **kwargs)
-        self.logp = logp
-        if type(self.logp) == types.MethodType:
-            if PLATFORM != "linux":
-                warnings.warn(
-                    "You are passing a bound method as logp for DensityDist, this can lead to "
-                    "errors when sampling on platforms other than Linux. Consider using a "
-                    "plain function instead, or subclass Distribution."
-                )
-            elif type(multiprocessing.get_context()) != multiprocessing.context.ForkContext:
-                warnings.warn(
-                    "You are passing a bound method as logp for DensityDist, this can lead to "
-                    "errors when sampling when multiprocessing cannot rely on forking. Consider using a "
-                    "plain function instead, or subclass Distribution."
-                )
-        self.rand = random
-        self.wrap_random_with_dist_shape = wrap_random_with_dist_shape
-        self.check_shape_in_random = check_shape_in_random
+        else:
+            dtype = cls.rv_op.dtype
+        ndim_supp = cls.rv_op.ndim_supp
+        if not hasattr(output.tag, "test_value"):
+            size = to_tuple(kwargs.get("size", None)) + (1,) * ndim_supp
+            output.tag.test_value = np.zeros(size, dtype)
+        return output
+
+
+def default_not_implemented(rv_name, method_name):
+    if method_name == "random":
+        # This is a hack to catch the NotImplementedError when creating the RV without random
+        # If the message starts with "Cannot sample from", then it uses the test_value as
+        # the initial_val.
+        message = (
+            f"Cannot sample from the DensityDist '{rv_name}' because the {method_name} "
+            "keyword argument was not provided when the distribution was "
+            f"but this method had not been provided when the distribution was "
+            f"constructed. Please re-build your model and provide a callable "
+            f"to '{rv_name}'s {method_name} keyword argument.\n"
+        )
+    else:
+        message = (
+            f"Attempted to run {method_name} on the DensityDist '{rv_name}', "
+            f"but this method had not been provided when the distribution was "
+            f"constructed. Please re-build your model and provide a callable "
+            f"to '{rv_name}'s {method_name} keyword argument.\n"
+        )
+
+    def func(*args, **kwargs):
+        raise NotImplementedError(message)
 
-    def _distr_parameters_for_repr(self):
-        return []
+    return func
diff --git a/pymc/gp/gp.py b/pymc/gp/gp.py
index b1a1dde678..9f913474d4 100644
--- a/pymc/gp/gp.py
+++ b/pymc/gp/gp.py
@@ -12,7 +12,6 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-import functools
 import warnings
 
 import aesara.tensor as at
@@ -730,12 +729,31 @@ def marginal_likelihood(self, name, X, Xu, y, noise=None, is_observed=True, **kw
                 )
         else:
             self.sigma = noise
-        logp = functools.partial(self._build_marginal_likelihood_logp, X=X, Xu=Xu, sigma=noise)
         if is_observed:
-            return pm.DensityDist(name, logp, observed=y, **kwargs)
+            return pm.DensityDist(
+                name,
+                X,
+                Xu,
+                self.sigma,
+                logp=self._build_marginal_likelihood_logp,
+                observed=y,
+                ndims_params=[2, 2, 0],
+                size=X.shape[0],
+                **kwargs,
+            )
         else:
-            shape = infer_shape(X, kwargs.pop("shape", None))
-            return pm.DensityDist(name, logp, size=shape, **kwargs)
+            return pm.DensityDist(
+                name,
+                X,
+                Xu,
+                self.sigma,
+                logp=self._build_marginal_likelihood_logp,
+                observed=y,
+                ndims_params=[2, 2, 0],
+                # ndim_supp=1,
+                size=X.shape[0],
+                **kwargs,
+            )
 
     def _build_conditional(self, Xnew, pred_noise, diag, X, Xu, y, sigma, cov_total, mean_total):
         sigma2 = at.square(sigma)
diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py
index ba1aef3f15..a272a5658b 100644
--- a/pymc/tests/test_distributions.py
+++ b/pymc/tests/test_distributions.py
@@ -120,6 +120,7 @@ def polyagamma_cdf(*args, **kwargs):
     logpt,
     logpt_sum,
 )
+from pymc.distributions.shape_utils import to_tuple
 from pymc.math import kronecker
 from pymc.model import Deterministic, Model, Point, Potential
 from pymc.tests.helpers import select_by_precision
@@ -1580,7 +1581,7 @@ def test_binomial(self):
             {"n": NatSmall, "p": Unit},
         )
 
-    @pytest.mark.xfail(reason="checkd tests has not been refactored")
+    @pytest.mark.xfail(reason="checkd tests have not been refactored")
     @pytest.mark.skipif(condition=(aesara.config.floatX == "float32"), reason="Fails on float32")
     def test_beta_binomial_distribution(self):
         self.checkd(
@@ -2448,7 +2449,7 @@ def test_orderedprobit(self, n):
             lambda value, eta, cutpoints: orderedprobit_logpdf(value, eta, cutpoints),
         )
 
-    @pytest.mark.xfail(reason="DensityDist no longer supported")
+    @pytest.mark.xfail(reason="checkd tests have not been refactored")
     def test_densitydist(self):
         def logp(x):
             return -log(2 * 0.5) - abs(x - 0.5) / 0.5
@@ -3227,14 +3228,13 @@ def test_issue_4499(self):
         assert_almost_equal(m.logp({"x": np.ones(10)}), 0 * 10)
 
 
-@pytest.mark.xfail(reason="DensityDist no longer supported")
 def test_serialize_density_dist():
     def func(x):
         return -2 * (x ** 2).sum()
 
     with pm.Model():
         pm.Normal("x")
-        y = pm.DensityDist("y", func)
+        y = pm.DensityDist("y", logp=func)
         pm.sample(draws=5, tune=1, mp_ctx="spawn")
 
     import cloudpickle
@@ -3287,3 +3287,34 @@ def test_logp_gives_migration_instructions(method, newcode):
         with pytest.raises(AttributeError, match=rf"use `{newcode}`"):
             f()
     pass
+
+
+def test_density_dist_old_api_error():
+    with pm.Model():
+        with pytest.raises(
+            TypeError, match="The DensityDist API has changed, you are using the old API"
+        ):
+            pm.DensityDist("a", lambda x: x)
+
+
+@pytest.mark.parametrize("size", [None, (), (2,)], ids=str)
+def test_density_dist_multivariate_logp(size):
+    supp_shape = 5
+    with pm.Model() as model:
+
+        def logp(value, mu):
+            return pm.MvNormal.logp(value, mu, at.eye(mu.shape[0]))
+
+        mu = pm.Normal("mu", size=supp_shape)
+        a = pm.DensityDist("a", mu, logp=logp, ndims_params=[1], ndim_supp=1, size=size)
+    mu_val = np.random.normal(loc=0, scale=1, size=supp_shape).astype(aesara.config.floatX)
+    a_val = np.random.normal(loc=mu_val, scale=1, size=to_tuple(size) + (supp_shape,)).astype(
+        aesara.config.floatX
+    )
+    log_densityt = pm.logpt(a, a.tag.value_var, sum=False)
+    assert (
+        log_densityt.eval(
+            {a.tag.value_var: a_val, mu.tag.value_var: mu_val},
+        ).shape
+        == to_tuple(size)
+    )
diff --git a/pymc/tests/test_distributions_random.py b/pymc/tests/test_distributions_random.py
index a8c8e7cf45..b6990c6a38 100644
--- a/pymc/tests/test_distributions_random.py
+++ b/pymc/tests/test_distributions_random.py
@@ -1838,111 +1838,54 @@ def test_mixture_random_shape_fast():
     assert rand3.shape == (100, 20)
 
 
-@pytest.mark.xfail(reason="This distribution has not been refactored for v4")
 class TestDensityDist:
-    @pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
-    def test_density_dist_with_random_sampleable(self, shape):
+    @pytest.mark.parametrize("size", [(), (3,), (3, 2)], ids=str)
+    def test_density_dist_with_random(self, size):
         with pm.Model() as model:
             mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1, shape=shape)
             obs = pm.DensityDist(
                 "density_dist",
-                normal_dist.logp,
-                observed=np.random.randn(100, *shape),
-                shape=shape,
-                random=normal_dist.random,
+                mu,
+                random=lambda mu, rng=None, size=None: rng.normal(loc=mu, scale=1, size=size),
+                observed=np.random.randn(100, *size),
+                size=size,
             )
-            idata = pm.sample(100, cores=1)
 
-        samples = 500
-        size = 100
-        ppc = pm.sample_posterior_predictive(idata, samples=samples, model=model, size=size)
-        assert ppc["density_dist"].shape == (samples, size) + obs.distribution.shape
+        assert obs.eval().shape == (100,) + size
 
-    @pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
-    def test_density_dist_with_random_sampleable_failure(self, shape):
+    def test_density_dist_without_random(self):
         with pm.Model() as model:
             mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1, shape=shape)
             pm.DensityDist(
                 "density_dist",
-                normal_dist.logp,
-                observed=np.random.randn(100, *shape),
-                shape=shape,
-                random=normal_dist.random,
-                wrap_random_with_dist_shape=False,
-            )
-            idata = pm.sample(100, cores=1)
-
-        samples = 500
-        with pytest.raises(RuntimeError):
-            pm.sample_posterior_predictive(idata, samples=samples, model=model, size=100)
-
-    @pytest.mark.parametrize("shape", [(), (3,), (3, 2)], ids=str)
-    def test_density_dist_with_random_sampleable_hidden_error(self, shape):
-        with pm.Model() as model:
-            mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1, shape=shape)
-            obs = pm.DensityDist(
-                "density_dist",
-                normal_dist.logp,
-                observed=np.random.randn(100, *shape),
-                shape=shape,
-                random=normal_dist.random,
-                wrap_random_with_dist_shape=False,
-                check_shape_in_random=False,
-            )
-            idata = pm.sample(100, cores=1)
-
-        samples = 500
-        ppc = pm.sample_posterior_predictive(idata, samples=samples, model=model)
-        assert len(ppc["density_dist"]) == samples
-        assert ((samples,) + obs.distribution.shape) != ppc["density_dist"].shape
-
-    def test_density_dist_with_random_sampleable_handcrafted_success(self):
-        with pm.Model() as model:
-            mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1)
-            rvs = pm.Normal.dist(mu, 1, shape=100).random
-            obs = pm.DensityDist(
-                "density_dist",
-                normal_dist.logp,
+                mu,
+                logp=lambda value, mu: pm.Normal.logp(value, mu, 1),
                 observed=np.random.randn(100),
-                random=rvs,
-                wrap_random_with_dist_shape=False,
             )
             idata = pm.sample(100, cores=1)
 
         samples = 500
-        size = 100
-        ppc = pm.sample_posterior_predictive(idata, samples=samples, model=model, size=size)
-        assert ppc["density_dist"].shape == (samples, size) + obs.distribution.shape
+        with pytest.raises(NotImplementedError):
+            pm.sample_posterior_predictive(idata, samples=samples, model=model, size=100)
 
-    @pytest.mark.xfail
-    def test_density_dist_with_random_sampleable_handcrafted_success_fast(self):
+    @pytest.mark.parametrize("size", [(), (3,), (3, 2)], ids=str)
+    def test_density_dist_with_random_multivariate(self, size):
+        supp_shape = 5
         with pm.Model() as model:
-            mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1)
-            rvs = pm.Normal.dist(mu, 1, shape=100).random
+            mu = pm.Normal("mu", 0, 1, size=supp_shape)
             obs = pm.DensityDist(
                 "density_dist",
-                normal_dist.logp,
-                observed=np.random.randn(100),
-                random=rvs,
-                wrap_random_with_dist_shape=False,
+                mu,
+                random=lambda mu, rng=None, size=None: rng.multivariate_normal(
+                    mean=mu, cov=np.eye(len(mu)), size=size
+                ),
+                observed=np.random.randn(100, *size, supp_shape),
+                size=size,
+                ndims_params=[1],
+                ndim_supp=1,
             )
-            pm.sample(100, cores=1)
 
-    def test_density_dist_without_random_not_sampleable(self):
-        with pm.Model() as model:
-            mu = pm.Normal("mu", 0, 1)
-            normal_dist = pm.Normal.dist(mu, 1)
-            pm.DensityDist("density_dist", normal_dist.logp, observed=np.random.randn(100))
-            idata = pm.sample(100, cores=1)
-
-        samples = 500
-        with pytest.raises(ValueError):
-            pm.sample_posterior_predictive(idata, samples=samples, model=model, size=100)
+        assert obs.eval().shape == (100,) + size + (supp_shape,)
 
 
 class TestNestedRandom(SeededTest):
diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py
index 5ae4d25b25..424e87e5e8 100644
--- a/pymc/tests/test_gp.py
+++ b/pymc/tests/test_gp.py
@@ -915,7 +915,7 @@ def testAdditiveMarginal(self):
         fp = np.random.randn(self.Xnew.shape[0])
         npt.assert_allclose(fp1.logp({"fp1": fp}), fp2.logp({"fp2": fp}), atol=0, rtol=1e-2)
 
-    @pytest.mark.xfail(reason="DensityDist was not yet refactored")
+    @pytest.mark.xfail(reason="MarginalSparse must be refactored to not use DensityDist")
     @pytest.mark.parametrize("approx", ["FITC", "VFE", "DTC"])
     def testAdditiveMarginalSparse(self, approx):
         Xu = np.random.randn(10, 3)
@@ -946,8 +946,10 @@ def testAdditiveMarginalSparse(self, approx):
         with model2:
             fp2 = gptot.conditional("fp2", self.Xnew)
 
-        fp = np.random.randn(self.Xnew.shape[0])
-        npt.assert_allclose(fp1.logp({"fp1": fp}), fp2.logp({"fp2": fp}), atol=0, rtol=1e-2)
+        fp = np.random.randn(self.Xnew.shape[0], 1)
+        npt.assert_allclose(
+            pm.logp(fp1, fp1).eval({fp1: fp}), pm.logp(fp2, fp2).eval({fp2: fp}), atol=0, rtol=1e-2
+        )
 
     @pytest.mark.xfail(reason="MvNormal was not yet refactored")
     def testAdditiveLatent(self):
diff --git a/pymc/tests/test_idata_conversion.py b/pymc/tests/test_idata_conversion.py
index a4dfffede2..93b38c0ccf 100644
--- a/pymc/tests/test_idata_conversion.py
+++ b/pymc/tests/test_idata_conversion.py
@@ -366,12 +366,15 @@ def test_multiple_observed_rv(self, log_likelihood):
         fails = check_multiple_attrs(test_dict, inference_data)
         assert not fails
 
-    @pytest.mark.xfail(reason="DensityDist not refactored for v4")
+    @pytest.mark.xfail(reason="MultiObservedRV is no longer used in v4")
     def test_multiple_observed_rv_without_observations(self):
         with pm.Model():
             mu = pm.Normal("mu")
             x = pm.DensityDist(  # pylint: disable=unused-variable
-                "x", logpt(pm.Normal.dist(mu, 1.0)), observed={"value": 0.1}
+                "x",
+                mu,
+                logp=lambda value, mu: pm.Normal.logp(value, mu, 1),
+                observed=0.1,
             )
             inference_data = pm.sample(100, chains=2, return_inferencedata=True)
         test_dict = {
@@ -384,7 +387,7 @@ def test_multiple_observed_rv_without_observations(self):
         assert not fails
         assert inference_data.observed_data.value.dtype.kind == "f"
 
-    @pytest.mark.xfail(reason="DensityDist not refactored for v4")
+    @pytest.mark.xfail(reason="MultiObservedRV is no longer used in v4")
     @pytest.mark.parametrize("multiobs", (True, False))
     def test_multiobservedrv_to_observed_data(self, multiobs):
         # fake regression data, with weights (W)
@@ -400,12 +403,13 @@ def test_multiobservedrv_to_observed_data(self, multiobs):
             b = pm.Normal("b", 0, 10)
             mu = a + b * X
             sigma = pm.HalfNormal("sigma", 1)
+            w = W
 
-            def weighted_normal(y, w):
-                return w * logpt(pm.Normal.dist(mu=mu, sd=sigma), y)
+            def weighted_normal(value, mu, sigma, w):
+                return w * pm.Normal.logp(value, mu, sigma)
 
             y_logp = pm.DensityDist(  # pylint: disable=unused-variable
-                "y_logp", weighted_normal, observed={"y": Y, "w": W}
+                "y_logp", mu, sigma, w, logp=weighted_normal, observed=Y, size=N
             )
             idata = pm.sample(
                 20, tune=20, return_inferencedata=True, idata_kwargs={"density_dist_obs": multiobs}
diff --git a/pymc/tests/test_model.py b/pymc/tests/test_model.py
index 61a23fc97c..c005f29990 100644
--- a/pymc/tests/test_model.py
+++ b/pymc/tests/test_model.py
@@ -346,7 +346,6 @@ def test_aesara_switch_broadcast_edge_cases_2(self):
         npt.assert_allclose(m.dlogp([m.rvs_to_values[mu]])({"mu": 0}), 2.499424682024436, rtol=1e-5)
 
 
-@pytest.mark.xfail(reason="DensityDist not refactored for v4")
 def test_multiple_observed_rv():
     "Test previously buggy multi-observed RV comparison code."
     y1_data = np.random.randn(10)
@@ -354,7 +353,7 @@ def test_multiple_observed_rv():
     with pm.Model() as model:
         mu = pm.Normal("mu")
         x = pm.DensityDist(  # pylint: disable=unused-variable
-            "x", pm.Normal.dist(mu, 1.0).logp, observed={"value": 0.1}
+            "x", mu, logp=lambda value, mu: pm.Normal.logp(value, mu, 1.0), observed=0.1
         )
     assert not model["x"] == model["mu"]
     assert model["x"] == model["x"]
diff --git a/pymc/tests/test_moment.py b/pymc/tests/test_moment.py
new file mode 100644
index 0000000000..ec8fd9450d
--- /dev/null
+++ b/pymc/tests/test_moment.py
@@ -0,0 +1,38 @@
+import aesara
+import numpy as np
+import pytest
+
+from aesara import tensor as at
+
+import pymc as pm
+
+from pymc.distributions.distribution import get_moment
+from pymc.distributions.shape_utils import to_tuple
+
+
+@pytest.mark.parametrize("size", [None, (), (2,), (3, 2)], ids=str)
+def test_density_dist_moment_scalar(size):
+    def moment(rv, size, mu):
+        return (at.ones(size) * mu).astype(rv.dtype)
+
+    mu_val = np.array(np.random.normal(loc=2, scale=1)).astype(aesara.config.floatX)
+    with pm.Model():
+        mu = pm.Normal("mu")
+        a = pm.DensityDist("a", mu, get_moment=moment, size=size)
+    evaled_moment = get_moment(a).eval({mu: mu_val})
+    assert evaled_moment.shape == to_tuple(size)
+    assert np.all(evaled_moment == mu_val)
+
+
+@pytest.mark.parametrize("size", [None, (), (2,), (3, 2)], ids=str)
+def test_density_dist_moment_multivariate(size):
+    def moment(rv, size, mu):
+        return (at.ones(size)[..., None] * mu).astype(rv.dtype)
+
+    mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX)
+    with pm.Model():
+        mu = pm.Normal("mu", size=5)
+        a = pm.DensityDist("a", mu, get_moment=moment, ndims_params=[1], ndim_supp=1, size=size)
+    evaled_moment = get_moment(a).eval({mu: mu_val})
+    assert evaled_moment.shape == to_tuple(size) + (5,)
+    assert np.all(evaled_moment == mu_val)
diff --git a/pymc/tests/test_parallel_sampling.py b/pymc/tests/test_parallel_sampling.py
index bdbfc9854f..645103d985 100644
--- a/pymc/tests/test_parallel_sampling.py
+++ b/pymc/tests/test_parallel_sampling.py
@@ -186,7 +186,6 @@ def test_iterator():
             pass
 
 
-@pytest.mark.xfail(reason="DensityDist was not yet refactored for v4")
 def test_spawn_densitydist_function():
     with pm.Model() as model:
         mu = pm.Normal("mu", 0, 1)
@@ -194,38 +193,19 @@ def test_spawn_densitydist_function():
         def func(x):
             return -2 * (x ** 2).sum()
 
-        obs = pm.DensityDist("density_dist", func, observed=np.random.randn(100))
+        obs = pm.DensityDist("density_dist", logp=func, observed=np.random.randn(100))
         pm.sample(draws=10, tune=10, step=pm.Metropolis(), cores=2, mp_ctx="spawn")
 
 
-@pytest.mark.xfail(reason="DensityDist was not yet refactored for v4")
 def test_spawn_densitydist_bound_method():
+    N = 100
     with pm.Model() as model:
         mu = pm.Normal("mu", 0, 1)
-        normal_dist = pm.Normal.dist(mu, 1)
-        logp = lambda x: pm.logp(normal_dist, x, transformed=False)
-        obs = pm.DensityDist("density_dist", logp, observed=np.random.randn(100))
-        msg = "logp for DensityDist is a bound method, leading to RecursionError while serializing"
-        with pytest.raises(ValueError, match=msg):
-            pm.sample(draws=10, tune=10, step=pm.Metropolis(), cores=2, mp_ctx="spawn")
+        normal_dist = pm.Normal.dist(mu, 1, size=N)
 
+        def logp(x):
+            out = pm.logp(normal_dist, x, transformed=False)
+            return out
 
-@pytest.mark.xfail(reason="DensityDist was not yet refactored for v4")
-def test_spawn_densitydist_syswarning(monkeypatch):
-    monkeypatch.setattr("pymc.distributions.distribution.PLATFORM", "win32")
-    with pm.Model() as model:
-        mu = pm.Normal("mu", 0, 1)
-        normal_dist = pm.Normal.dist(mu, 1)
-        with pytest.warns(UserWarning, match="errors when sampling on platforms"):
-            obs = pm.DensityDist("density_dist", normal_dist.logp, observed=np.random.randn(100))
-
-
-@pytest.mark.xfail(reason="DensityDist was not yet refactored for v4")
-def test_spawn_densitydist_mpctxwarning(monkeypatch):
-    ctx = multiprocessing.get_context("spawn")
-    monkeypatch.setattr(multiprocessing, "get_context", lambda: ctx)
-    with pm.Model() as model:
-        mu = pm.Normal("mu", 0, 1)
-        normal_dist = pm.Normal.dist(mu, 1)
-        with pytest.warns(UserWarning, match="errors when sampling when multiprocessing"):
-            obs = pm.DensityDist("density_dist", normal_dist.logp, observed=np.random.randn(100))
+        obs = pm.DensityDist("density_dist", logp=logp, observed=np.random.randn(N), size=N)
+        pm.sample(draws=10, tune=10, step=pm.Metropolis(), cores=2, mp_ctx="spawn")
diff --git a/pymc/tests/test_sampling.py b/pymc/tests/test_sampling.py
index 7d929a428b..ad65e4c164 100644
--- a/pymc/tests/test_sampling.py
+++ b/pymc/tests/test_sampling.py
@@ -1043,7 +1043,6 @@ def test_shared(self):
         assert gen2["y"].shape == (draws, n2)
         assert gen2["o"].shape == (draws, n2)
 
-    @pytest.mark.xfail(reason="DensityDist not refactored for v4")
     def test_density_dist(self):
         obs = np.random.normal(-1, 0.1, size=10)
         with pm.Model():
@@ -1051,8 +1050,9 @@ def test_density_dist(self):
             sd = pm.Gamma("sd", 1, 2)
             a = pm.DensityDist(
                 "a",
-                pm.Normal.dist(mu, sd).logp,
-                random=pm.Normal.dist(mu, sd).random,
+                mu,
+                sd,
+                random=lambda mu, sd, rng=None, size=None: rng.normal(loc=mu, scale=sd, size=size),
                 observed=obs,
             )
             prior = pm.sample_prior_predictive()


Code review comments:
diff_hunk: @@ -58,16 +58,16 @@ An exponential survival function, where :math:`c=0` denotes failure (or non-surv
     f(c, t) = \left\{ \begin{array}{l} \exp(-\lambda t), \text{if c=1} \\
                \lambda \exp(-\lambda t), \text{if c=0}  \end{array} \right.
 
-Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as an argument to the ``DensityDist`` function, which creates an instance of a PyMC distribution with the custom function as its log-probability.
+Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as a keyword argument to the ``DensityDist`` function, which creates an instance of a PyMC3 distribution with the custom function as its log-probability.
 
 For the exponential survival function, this is:
 
 ::
 
-    def logp(failure, value):
-        return (failure * log(λ) - λ * value).sum()
+    def logp(value, t, λ):
path: docs/source/Probability_Distributions.rst
text: @user1:
Earlier this year I put a sigma character in a signature, because I thought it was neat.

Two months later I took it back out, because it was so annoying to type Greek letters outside of Jupyter notebooks.

"lam" is the alternative we typically use in other places. 

@author:
Thanks for the comments @user1, I’ll edit the code as soon as possible.

diff_hunk: @@ -387,96 +390,241 @@ class NoDistribution(Distribution):
     """
 
 
-class DensityDist(Distribution):
-    """Distribution based on a given log density function.
+class DensityDistRV(RandomVariable):
+    """
+    Base class for DensityDistRV
+
+    This should be subclassed when defining custom DensityDist objects.
+    """
+
+    name = "DensityDistRV"
+    _print_name = ("DensityDist", "\\operatorname{DensityDist}")
+
+    @classmethod
+    def rng_fn(cls, rng, *args):
+        args = list(args)
+        size = args.pop(-1)
+        return cls._random_fn(*args, rng=rng, size=size)
 
-    A distribution with the passed log density function is created.
-    Requires a custom random function passed as kwarg `random` to
-    enable prior or posterior predictive sampling.
 
+class DensityDist(NoDistribution):
+    """A distribution that can be used to wrap black-box log density functions.
+
+    Creates a Distribution and registers the supplied log density function to be used
+    for inference. It is also possible to supply a `random` method in order to be able
+    to sample from the prior or posterior predictive distributions.
     """
 
-    def __init__(
-        self,
-        logp,
-        shape=(),
-        dtype=None,
-        initval=0,
-        random=None,
-        wrap_random_with_dist_shape=True,
-        check_shape_in_random=True,
-        *args,
+    def __new__(
+        cls,
+        name: str,
+        *dist_params,
+        logp: Optional[Callable] = None,
+        logcdf: Optional[Callable] = None,
+        random: Optional[Callable] = None,
+        get_moment: Optional[Callable] = None,
+        ndim_supp: int = 0,
+        ndims_params: Optional[Sequence[int]] = None,
+        dtype: str = "floatX",
         **kwargs,
     ):
         """
         Parameters
         ----------
-
-        logp: callable
-            A callable that has the following signature ``logp(value)`` and
-            returns an Aesara tensor that represents the distribution's log
-            probability density.
-        shape: tuple (Optional): defaults to `()`
-            The shape of the distribution. The default value indicates a scalar.
-            If the distribution is *not* scalar-valued, the programmer should pass
-            a value here.
-        dtype: None, str (Optional)
-            The dtype of the distribution.
-        initval: number or array (Optional)
-            The ``initval`` of the RV's tensor that follow the ``DensityDist``
-            distribution.
-        args, kwargs: (Optional)
-            These are passed to the parent class' ``__init__``.
+        name : str
+        dist_params : Tuple
+            A sequence of the distribution's parameter. These will be converted into
+            aesara tensors internally. These parameters could be other ``RandomVariable``
path: pymc/distributions/distribution.py
text: @user1:
Aesara capitalization

diff_hunk: @@ -58,16 +58,16 @@ An exponential survival function, where :math:`c=0` denotes failure (or non-surv
     f(c, t) = \left\{ \begin{array}{l} \exp(-\lambda t), \text{if c=1} \\
                \lambda \exp(-\lambda t), \text{if c=0}  \end{array} \right.
 
-Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as an argument to the ``DensityDist`` function, which creates an instance of a PyMC distribution with the custom function as its log-probability.
+Such a function can be implemented as a PyMC distribution by writing a function that specifies the log-probability, then passing that function as a keyword argument to the ``DensityDist`` function, which creates an instance of a PyMC3 distribution with the custom function as its log-probability.
path: docs/source/Probability_Distributions.rst
text: @user1:
PyMC3 - > PyMC