Repository: pymc-devs/pymc

Issue:
BUG: object of type 'NoneType' has no len() when computing log likelihood
### Describe the issue:

I have see this error when computing the log likelihood in hierarchical models. If you do not calculate the log likelihood the model samples just fine and if your model is not hierarchical you can calculate the log likelihood. I have seen this issue in my models but you can reproduce it using the radon example. When computing the log likelihood with either 

```
with partial_pooling:
    partial_pooling_trace = pm.sample(tune=2000, random_seed=RANDOM_SEED, idata_kwargs=dict(log_likelihood=True))
```

or with

```
with partial_pooling:
    pm.compute_log_likelihood(partial_pooling_trace)
```

The code below is copied out of the simplest partial pooling model from the radon example.

### Reproduceable code example:

```python
import os
os.environ["PATH"] += os.pathsep + 'C:/Program Files/Graphviz/bin'
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import seaborn as sns
import xarray as xr

warnings.filterwarnings("ignore", module="scipy")
RANDOM_SEED = 8924

try:
    srrs2 = pd.read_csv(os.path.join("..", "data", "srrs2.dat"))
except FileNotFoundError:
    srrs2 = pd.read_csv(pm.get_data("srrs2.dat"))

srrs2.columns = srrs2.columns.map(str.strip)
srrs_mn = srrs2[srrs2.state == "MN"].copy()

try:
    cty = pd.read_csv(os.path.join("..", "data", "cty.dat"))
except FileNotFoundError:
    cty = pd.read_csv(pm.get_data("cty.dat"))

srrs_mn["fips"] = srrs_mn.stfips * 1000 + srrs_mn.cntyfips
cty_mn = cty[cty.st == "MN"].copy()
cty_mn["fips"] = 1000 * cty_mn.stfips + cty_mn.ctfips

srrs_mn = srrs_mn.merge(cty_mn[["fips", "Uppm"]], on="fips")
srrs_mn = srrs_mn.drop_duplicates(subset="idnum")
u = np.log(srrs_mn.Uppm).unique()

n = len(srrs_mn)

srrs_mn.county = srrs_mn.county.map(str.strip)
county, mn_counties = srrs_mn.county.factorize()
srrs_mn["county_code"] = county
radon = srrs_mn.activity
srrs_mn["log_radon"] = log_radon = np.log(radon + 0.1).values
floor_measure = srrs_mn.floor.values

with pm.Model(coords=coords) as partial_pooling:
    county_idx = pm.MutableData("county_idx", county, dims="obs_id")

    # Priors
    mu_a = pm.Normal("mu_a", mu=0.0, sigma=10)
    sigma_a = pm.Exponential("sigma_a", 1)

    # Random intercepts
    alpha = pm.Normal("alpha", mu=mu_a, sigma=sigma_a, dims="county")

    # Model error
    sigma_y = pm.Exponential("sigma_y", 1)

    # Expected value
    y_hat = alpha[county_idx]

    # Data likelihood
    y_like = pm.Normal("y_like", mu=y_hat, sigma=sigma_y, observed=log_radon, dims="obs_id")

with partial_pooling:
    partial_pooling_trace = pm.sample(tune=2000, random_seed=RANDOM_SEED, idata_kwargs=dict(log_likelihood=True))
```


### Error message:

```shell
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[27], line 2
      1 with partial_pooling:
----> 2     partial_pooling_trace = pm.sample(tune=2000, random_seed=RANDOM_SEED, idata_kwargs=dict(log_likelihood=True))

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\pymc\sampling\mcmc.py:791, in sample(draws, tune, chains, cores, random_seed, progressbar, step, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, model, **kwargs)
    787 t_sampling = time.time() - t_start
    789 # Packaging, validating and returning the result was extracted
    790 # into a function to make it easier to test and refactor.
--> 791 return _sample_return(
    792     run=run,
    793     traces=traces,
    794     tune=tune,
    795     t_sampling=t_sampling,
    796     discard_tuned_samples=discard_tuned_samples,
    797     compute_convergence_checks=compute_convergence_checks,
    798     return_inferencedata=return_inferencedata,
    799     keep_warning_stat=keep_warning_stat,
    800     idata_kwargs=idata_kwargs or {},
    801     model=model,
    802 )

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\pymc\sampling\mcmc.py:859, in _sample_return(run, traces, tune, t_sampling, discard_tuned_samples, compute_convergence_checks, return_inferencedata, keep_warning_stat, idata_kwargs, model)
    857 ikwargs: Dict[str, Any] = dict(model=model, save_warmup=not discard_tuned_samples)
    858 ikwargs.update(idata_kwargs)
--> 859 idata = pm.to_inference_data(mtrace, **ikwargs)
    861 if compute_convergence_checks:
    862     warns = run_convergence_checks(idata, model)

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\pymc\backends\arviz.py:519, in to_inference_data(trace, prior, posterior_predictive, log_likelihood, coords, dims, sample_dims, model, save_warmup, include_transformed)
    505 if isinstance(trace, InferenceData):
    506     return trace
    508 return InferenceDataConverter(
    509     trace=trace,
    510     prior=prior,
    511     posterior_predictive=posterior_predictive,
    512     log_likelihood=log_likelihood,
    513     coords=coords,
    514     dims=dims,
    515     sample_dims=sample_dims,
    516     model=model,
    517     save_warmup=save_warmup,
    518     include_transformed=include_transformed,
--> 519 ).to_inference_data()

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\pymc\backends\arviz.py:442, in InferenceDataConverter.to_inference_data(self)
    439 if self.log_likelihood:
    440     from pymc.stats.log_likelihood import compute_log_likelihood
--> 442     idata = compute_log_likelihood(
    443         idata,
    444         var_names=None if self.log_likelihood is True else self.log_likelihood,
    445         extend_inferencedata=True,
    446         model=self.model,
    447         sample_dims=self.sample_dims,
    448         progressbar=False,
    449     )
    450 return idata

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\pymc\stats\log_likelihood.py:116, in compute_log_likelihood(idata, var_names, extend_inferencedata, model, sample_dims, progressbar)
    111 for key, array in loglike_trace.items():
    112     loglike_trace[key] = array.reshape(
    113         (*[len(coord) for coord in stacked_dims.values()], *array.shape[1:])
    114     )
--> 116 loglike_dataset = dict_to_dataset(
    117     loglike_trace,
    118     library=pymc,
    119     dims={dname: list(dvals) for dname, dvals in model.named_vars_to_dims.items()},
    120     coords={
    121         cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
    122         for cname, cvals in model.coords.items()
    123     },
    124     default_dims=list(sample_dims),
    125     skip_event_dims=True,
    126 )
    128 if extend_inferencedata:
    129     idata.add_groups(dict(log_likelihood=loglike_dataset))

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\arviz\data\base.py:306, in dict_to_dataset(data, attrs, library, coords, dims, default_dims, index_origin, skip_event_dims)
    303 if dims is None:
    304     dims = {}
--> 306 data_vars = {
    307     key: numpy_to_data_array(
    308         values,
    309         var_name=key,
    310         coords=coords,
    311         dims=dims.get(key),
    312         default_dims=default_dims,
    313         index_origin=index_origin,
    314         skip_event_dims=skip_event_dims,
    315     )
    316     for key, values in data.items()
    317 }
    318 return xr.Dataset(data_vars=data_vars, attrs=make_attrs(attrs=attrs, library=library))

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\arviz\data\base.py:307, in <dictcomp>(.0)
    303 if dims is None:
    304     dims = {}
    306 data_vars = {
--> 307     key: numpy_to_data_array(
    308         values,
    309         var_name=key,
    310         coords=coords,
    311         dims=dims.get(key),
    312         default_dims=default_dims,
    313         index_origin=index_origin,
    314         skip_event_dims=skip_event_dims,
    315     )
    316     for key, values in data.items()
    317 }
    318 return xr.Dataset(data_vars=data_vars, attrs=make_attrs(attrs=attrs, library=library))

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\arviz\data\base.py:231, in numpy_to_data_array(ary, var_name, coords, dims, default_dims, index_origin, skip_event_dims)
    228 else:
    229     ary = utils.one_de(ary)
--> 231 dims, coords = generate_dims_coords(
    232     ary.shape[len(default_dims) :],
    233     var_name,
    234     dims=dims,
    235     coords=coords,
    236     default_dims=default_dims,
    237     index_origin=index_origin,
    238     skip_event_dims=skip_event_dims,
    239 )
    241 # reversed order for default dims: 'chain', 'draw'
    242 if "draw" not in dims and "draw" in default_dims:

File ~\AppData\Local\Programs\Python\Python310\lib\site-packages\arviz\data\base.py:146, in generate_dims_coords(shape, var_name, dims, coords, default_dims, index_origin, skip_event_dims)
    143 if skip_event_dims:
    144     # this is needed in case the reduction keeps the dimension with size 1
    145     for i, (dim, dim_size) in enumerate(zip(dims, shape)):
--> 146         if (dim in coords) and (dim_size != len(coords[dim])):
    147             dims = dims[:i]
    148             break

TypeError: object of type 'NoneType' has no len()
```


### PyMC version information:

pymc version 5.6.0
pytensor version 2.12.3

I also got this error using pymc version 5.5.0

### Context for the issue:

Users cannot calculate log likelihoods for hierarchical models which limits their ability to do model comparisons.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#6882](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (cab6332) into [main](https://app.codecov.io/gh/pymc-devs/pymc/commit/ddd1d4bf05a72895c67265f541585ae02bd338a3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (ddd1d4b) will **decrease** coverage by `51.11%`.
> The diff coverage is `n/a`.

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #6882       +/-   ##
===========================================
- Coverage   92.05%   40.95%   -51.11%     
===========================================
  Files          96       96               
  Lines       16448    16446        -2     
===========================================
- Hits        15142     6736     -8406     
- Misses       1306     9710     +8404     
```


| [Files Changed](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/stats/log\_likelihood.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGF0cy9sb2dfbGlrZWxpaG9vZC5weQ==) | `89.36% <ø> (-10.64%)` | :arrow_down: |

... and [77 files with indirect coverage changes](https://app.codecov.io/gh/pymc-devs/pymc/pull/6882/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

</details>
ricardoV94: Thanks @jaharvey8. Can you add a small test that was triggering the error before this fix?
jaharvey8: > 

Sure, I hate to ask an annoying question, but what's the best format for uploading a test? Or is there documentation somewhere on uploading tests?
ricardoV94: > Sure, I hate to ask an annoying question, but what's the best format for uploading a test? Or is there documentation somewhere on uploading tests?

Not annoying at all. There should some tests in `tests/stats/test_log_likelihood.py` (or something similar). You can add a separate test function for this bug there.

To run the tests locally check this guide: https://www.pymc.io/projects/docs/en/latest/contributing/running_the_test_suite.html

If you have any questions, just ask away!
ricardoV94: @jaharvey8 I refactored a bit the code to avoid duplication and added a test. Let me know if it looks good on your end.

Patch:
diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py
index bfd3008a3..39553f625 100644
--- a/pymc/backends/arviz.py
+++ b/pymc/backends/arviz.py
@@ -100,6 +100,18 @@ def find_constants(model: "Model") -> Dict[str, Var]:
     return constant_data
 
 
+def coords_and_dims_for_inferencedata(model: Model) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Parse PyMC model coords and dims format to one accepted by InferenceData."""
+    coords = {
+        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
+        for cname, cvals in model.coords.items()
+        if cvals is not None
+    }
+    dims = {dname: list(dvals) for dname, dvals in model.named_vars_to_dims.items()}
+
+    return coords, dims
+
+
 class _DefaultTrace:
     """
     Utility for collecting samples into a dictionary.
@@ -216,19 +228,11 @@ class InferenceDataConverter:  # pylint: disable=too-many-instance-attributes
                 " one of trace, prior, posterior_predictive or predictions."
             )
 
-        # Make coord types more rigid
-        untyped_coords: Dict[str, Optional[Sequence[Any]]] = {**self.model.coords}
-        if coords:
-            untyped_coords.update(coords)
-        self.coords = {
-            cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-            for cname, cvals in untyped_coords.items()
-            if cvals is not None
-        }
-
-        self.dims = {} if dims is None else dims
-        model_dims = {k: list(v) for k, v in self.model.named_vars_to_dims.items()}
-        self.dims = {**model_dims, **self.dims}
+        user_coords = {} if coords is None else coords
+        user_dims = {} if dims is None else dims
+        model_coords, model_dims = coords_and_dims_for_inferencedata(self.model)
+        self.coords = {**model_coords, **user_coords}
+        self.dims = {**model_dims, **user_dims}
 
         if sample_dims is None:
             sample_dims = ["chain", "draw"]
diff --git a/pymc/model/core.py b/pymc/model/core.py
index 2edcb1836..01a655679 100644
--- a/pymc/model/core.py
+++ b/pymc/model/core.py
@@ -942,7 +942,7 @@ class Model(WithMemoization, metaclass=ContextMeta):
         Entries in the tuples may be ``None``, if the RV dimension was not given a name.
         """
         warnings.warn(
-            "Model.RV_dims is deprecated. User Model.named_vars_to_dims instead.",
+            "Model.RV_dims is deprecated. Use Model.named_vars_to_dims instead.",
             FutureWarning,
         )
         return self.named_vars_to_dims
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
index 5016b0897..e0174c257 100644
--- a/pymc/sampling/jax.py
+++ b/pymc/sampling/jax.py
@@ -37,7 +37,11 @@ from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.type import RandomType
 
 from pymc import Model, modelcontext
-from pymc.backends.arviz import find_constants, find_observations
+from pymc.backends.arviz import (
+    coords_and_dims_for_inferencedata,
+    find_constants,
+    find_observations,
+)
 from pymc.distributions.multivariate import PosDefMatrix
 from pymc.initial_point import StartDict
 from pymc.logprob.utils import CheckParameterValue
@@ -392,17 +396,6 @@ def sample_blackjax_nuts(
 
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
-    coords = {
-        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-        for cname, cvals in model.coords.items()
-        if cvals is not None
-    }
-
-    dims = {
-        var_name: [dim for dim in dims if dim is not None]
-        for var_name, dims in model.named_vars_to_dims.items()
-    }
-
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
     tic1 = datetime.now()
@@ -485,7 +478,7 @@ def sample_blackjax_nuts(
         "sampling_time": (tic3 - tic2).total_seconds(),
     }
 
-    posterior = mcmc_samples
+    coords, dims = coords_and_dims_for_inferencedata(model)
     # Update 'coords' and 'dims' extracted from the model with user 'idata_kwargs'
     # and drop keys 'coords' and 'dims' from 'idata_kwargs' if present.
     _update_coords_and_dims(coords=coords, dims=dims, idata_kwargs=idata_kwargs)
@@ -500,7 +493,7 @@ def sample_blackjax_nuts(
         dims=dims,
         attrs=make_attrs(attrs, library=blackjax),
     )
-    az_trace = to_trace(posterior=posterior, **idata_kwargs)
+    az_trace = to_trace(posterior=mcmc_samples, **idata_kwargs)
 
     return az_trace
 
@@ -613,17 +606,6 @@ def sample_numpyro_nuts(
 
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
-    coords = {
-        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-        for cname, cvals in model.coords.items()
-        if cvals is not None
-    }
-
-    dims = {
-        var_name: [dim for dim in dims if dim is not None]
-        for var_name, dims in model.named_vars_to_dims.items()
-    }
-
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
     tic1 = datetime.now()
@@ -715,7 +697,7 @@ def sample_numpyro_nuts(
         "sampling_time": (tic3 - tic2).total_seconds(),
     }
 
-    posterior = mcmc_samples
+    coords, dims = coords_and_dims_for_inferencedata(model)
     # Update 'coords' and 'dims' extracted from the model with user 'idata_kwargs'
     # and drop keys 'coords' and 'dims' from 'idata_kwargs' if present.
     _update_coords_and_dims(coords=coords, dims=dims, idata_kwargs=idata_kwargs)
@@ -730,5 +712,5 @@ def sample_numpyro_nuts(
         dims=dims,
         attrs=make_attrs(attrs, library=numpyro),
     )
-    az_trace = to_trace(posterior=posterior, **idata_kwargs)
+    az_trace = to_trace(posterior=mcmc_samples, **idata_kwargs)
     return az_trace
diff --git a/pymc/stats/log_likelihood.py b/pymc/stats/log_likelihood.py
index 51bbd59d6..99f6fe536 100644
--- a/pymc/stats/log_likelihood.py
+++ b/pymc/stats/log_likelihood.py
@@ -13,14 +13,12 @@
 #   limitations under the License.
 from typing import Optional, Sequence, cast
 
-import numpy as np
-
 from arviz import InferenceData, dict_to_dataset
 from fastprogress import progress_bar
 
 import pymc
 
-from pymc.backends.arviz import _DefaultTrace
+from pymc.backends.arviz import _DefaultTrace, coords_and_dims_for_inferencedata
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import PointFunc
 from pymc.util import dataset_to_point_list
@@ -113,14 +111,12 @@ def compute_log_likelihood(
             (*[len(coord) for coord in stacked_dims.values()], *array.shape[1:])
         )
 
+    coords, dims = coords_and_dims_for_inferencedata(model)
     loglike_dataset = dict_to_dataset(
         loglike_trace,
         library=pymc,
-        dims={dname: list(dvals) for dname, dvals in model.named_vars_to_dims.items()},
-        coords={
-            cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-            for cname, cvals in model.coords.items()
-        },
+        dims=dims,
+        coords=coords,
         default_dims=list(sample_dims),
         skip_event_dims=True,
     )


Test patch:
diff --git a/tests/stats/test_log_likelihood.py b/tests/stats/test_log_likelihood.py
index e7201b761..fde4d69c3 100644
--- a/tests/stats/test_log_likelihood.py
+++ b/tests/stats/test_log_likelihood.py
@@ -15,7 +15,7 @@ import numpy as np
 import pytest
 import scipy.stats as st
 
-from arviz import InferenceData, dict_to_dataset
+from arviz import InferenceData, dict_to_dataset, from_dict
 
 from pymc.distributions import Dirichlet, Normal
 from pymc.distributions.transforms import log
@@ -117,3 +117,18 @@ class TestComputeLogLikelihood:
             idata = InferenceData(posterior=dict_to_dataset({"x": np.arange(100).reshape(4, 25)}))
             with pytest.raises(ValueError, match="var_names must refer to observed_RVs"):
                 compute_log_likelihood(idata, var_names=["x"])
+
+    def test_dims_without_coords(self):
+        # Issues #6820
+        with Model() as m:
+            x = Normal("x")
+            y = Normal("y", x, observed=[0, 0, 0], shape=(3,), dims="obs")
+
+            trace = from_dict({"x": [[0, 1]]})
+            llike = compute_log_likelihood(trace)
+
+        assert len(llike.log_likelihood["obs"]) == 3
+        np.testing.assert_allclose(
+            llike.log_likelihood["y"].values,
+            st.norm.logpdf([[[0, 0, 0], [1, 1, 1]]]),
+        )