Make VI work on v4

.github/workflows/pytest.yml

@@ -48,7 +48,6 @@ jobs:
             --ignore=pymc/tests/test_step.py
             --ignore=pymc/tests/test_tuning.py
             --ignore=pymc/tests/test_transforms.py
-            --ignore=pymc/tests/test_variational_inference.py
             --ignore=pymc/tests/test_sampling_jax.py
             --ignore=pymc/tests/test_dist_math.py
             --ignore=pymc/tests/test_minibatches.py
@@ -169,6 +168,7 @@ jobs:
             pymc/tests/test_distributions_random.py
             pymc/tests/test_distributions_moments.py
             pymc/tests/test_distributions_timeseries.py
+            pymc/tests/test_variational_inference.py
           - |
             pymc/tests/test_parallel_sampling.py
             pymc/tests/test_sampling.py

pymc/backends/arviz.py

@@ -478,6 +478,7 @@ def is_data(name, var) -> bool:
                 and var not in self.model.observed_RVs
                 and var not in self.model.free_RVs
                 and var not in self.model.potentials
+                and var not in self.model.value_vars
                 and (self.observations is None or name not in self.observations)
                 and isinstance(var, (Constant, SharedVariable))
             )

pymc/distributions/logprob.py

@@ -14,8 +14,9 @@
 
 from collections.abc import Mapping
 from functools import singledispatch
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Sequence, Union
 
+import aesara
 import aesara.tensor as at
 import numpy as np
 
@@ -43,15 +44,17 @@ def logp_transform(op: Op):
     return None
 
 
-def _get_scaling(total_size, shape, ndim):
+def _get_scaling(total_size: Optional[Union[int, Sequence[int]]], shape, ndim: int):
     """
-    Gets scaling constant for logp
+    Gets scaling constant for logp.
 
     Parameters
     ----------
-    total_size: int or list[int]
+    total_size: Optional[int|List[int]]
+        size of a fully observed data without minibatching,
+        `None` means data is fully observed
     shape: shape
-        shape to scale
+        shape of an observed data
     ndim: int
         ndim hint
 
@@ -60,7 +63,7 @@ def _get_scaling(total_size, shape, ndim):
     scalar
     """
     if total_size is None:
-        coef = floatX(1)
+        coef = 1.0
     elif isinstance(total_size, int):
         if ndim >= 1:
             denom = shape[0]
@@ -90,21 +93,23 @@ def _get_scaling(total_size, shape, ndim):
                 "number of scalings is bigger that ndim, got %r" % total_size
             )
         elif (len(begin) + len(end)) == 0:
-            return floatX(1)
+            coef = 1.0
         if len(end) > 0:
             shp_end = shape[-len(end) :]
         else:
             shp_end = np.asarray([])
         shp_begin = shape[: len(begin)]
-        begin_coef = [floatX(t) / shp_begin[i] for i, t in enumerate(begin) if t is not None]
-        end_coef = [floatX(t) / shp_end[i] for i, t in enumerate(end) if t is not None]
+        begin_coef = [
+            floatX(t) / floatX(shp_begin[i]) for i, t in enumerate(begin) if t is not None
+        ]
+        end_coef = [floatX(t) / floatX(shp_end[i]) for i, t in enumerate(end) if t is not None]
         coefs = begin_coef + end_coef
         coef = at.prod(coefs)
     else:
         raise TypeError(
             "Unrecognized `total_size` type, expected int or list of ints, got %r" % total_size
         )
-    return at.as_tensor(floatX(coef))
+    return at.as_tensor(coef, dtype=aesara.config.floatX)
 
 
 subtensor_types = (

pymc/model.py

@@ -58,6 +58,7 @@
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.data import GenTensorVariable, Minibatch
 from pymc.distributions import joint_logpt, logp_transform
+from pymc.distributions.logprob import _get_scaling
 from pymc.exceptions import ImputationWarning, SamplingError, ShapeError
 from pymc.initial_point import make_initial_point_fn
 from pymc.math import flatten_list
@@ -1238,6 +1239,7 @@ def register_rv(
         name = self.name_for(name)
         rv_var.name = name
         rv_var.tag.total_size = total_size
+        rv_var.tag.scaling = _get_scaling(total_size, shape=rv_var.shape, ndim=rv_var.ndim)
 
         # Associate previously unknown dimension names with
         # the length of the corresponding RV dimension.
@@ -1870,7 +1872,7 @@ def Potential(name, var, model=None):
     """
     model = modelcontext(model)
     var.name = model.name_for(name)
-    var.tag.scaling = None
+    var.tag.scaling = 1.0
     model.potentials.append(var)
     model.add_random_variable(var)
 

pymc/sampling.py

@@ -2385,7 +2385,7 @@ def init_nuts(
             progressbar=progressbar,
             obj_optimizer=pm.adagrad_window,
         )
-        initial_points = list(approx.sample(draws=chains))
+        initial_points = list(approx.sample(draws=chains, return_inferencedata=False))
         std_apoint = approx.std.eval()
         cov = std_apoint**2
         mean = approx.mean.get_value()
@@ -2402,7 +2402,7 @@ def init_nuts(
             progressbar=progressbar,
             obj_optimizer=pm.adagrad_window,
         )
-        initial_points = list(approx.sample(draws=chains))
+        initial_points = list(approx.sample(draws=chains, return_inferencedata=False))
         cov = approx.std.eval() ** 2
         potential = quadpotential.QuadPotentialDiag(cov)
     elif init == "advi_map":
@@ -2416,7 +2416,7 @@ def init_nuts(
             progressbar=progressbar,
             obj_optimizer=pm.adagrad_window,
         )
-        initial_points = list(approx.sample(draws=chains))
+        initial_points = list(approx.sample(draws=chains, return_inferencedata=False))
         cov = approx.std.eval() ** 2
         potential = quadpotential.QuadPotentialDiag(cov)
     elif init == "map":