BUG: Issues creating InferenceData if there is a coord of type Sequence[Sequence]

[covertg]

Describe the issue:

Hi, writing with a niche question. I tried to specify a model where I've flattened multiple coords into a "flattened" single coord, where each value is a particular combination of the original coords. E.g. instead of {"coord1": ["a", "b"], "coord2": [1, 2, 3]}, I set {"c": [("a", 1), ("a", 2), ("a", 3), ("b", 1), ("b", 2), ("b", 3)].

Is this supported input? Unfortunately, sampling fails in this case.

Reproduceable code example:

import pymc as pm
with pm.Model(coords={"c": [("a", 1), ("a", 3), ("b", 2)]}) as m:
    x = pm.Normal("x", mu=0, sigma=1, dims="c")
    # m.coords is ({'c': (('a', 1), ('a', 3), ('b', 2))}
    # x.type is TensorType(float64, (3,)))
    idata = pm.sample_prior_predictive()
    # also fails with,
    # idata = pm.sample()

Error message:

ValueError                                Traceback (most recent call last)
Cell In[23], line 2
      1 with m:
----> 2     idata = pm.sample_prior_predictive()

File ...python/site-packages/pymc/sampling/forward.py:440, in sample_prior_predictive(samples, model, var_names, random_seed, return_inferencedata, idata_kwargs, compile_kwargs)
    438 if idata_kwargs:
    439     ikwargs.update(idata_kwargs)
--> 440 return pm.to_inference_data(prior=prior, **ikwargs)

File ...python/site-packages/pymc/backends/arviz.py:496, in to_inference_data(trace, prior, posterior_predictive, log_likelihood, coords, dims, sample_dims, model, save_warmup, include_transformed)
    482 if isinstance(trace, InferenceData):
    483     return trace
    485 return InferenceDataConverter(
    486     trace=trace,
    487     prior=prior,
    488     posterior_predictive=posterior_predictive,
    489     log_likelihood=log_likelihood,
    490     coords=coords,
    491     dims=dims,
    492     sample_dims=sample_dims,
    493     model=model,
    494     save_warmup=save_warmup,
    495     include_transformed=include_transformed,
--> 496 ).to_inference_data()

File ...python/site-packages/pymc/backends/arviz.py:408, in InferenceDataConverter.to_inference_data(self)
    396 def to_inference_data(self):
    397     """Convert all available data to an InferenceData object.
    398 
    399     Note that if groups can not be created (e.g., there is no `trace`, so
    400     the `posterior` and `sample_stats` can not be extracted), then the InferenceData
    401     will not have those groups.
    402     """
    403     id_dict = {
    404         "posterior": self.posterior_to_xarray(),
    405         "sample_stats": self.sample_stats_to_xarray(),
    406         "posterior_predictive": self.posterior_predictive_to_xarray(),
    407         "predictions": self.predictions_to_xarray(),
--> 408         **self.priors_to_xarray(),
    409         "observed_data": self.observed_data_to_xarray(),
    410     }
    411     if self.predictions:
    412         id_dict["predictions_constant_data"] = self.constant_data_to_xarray()

File ...python/site-packages/pymc/backends/arviz.py:358, in InferenceDataConverter.priors_to_xarray(self)
    351 priors_dict = {}
    352 for group, var_names in zip(
    353     ("prior", "prior_predictive"), (prior_vars, prior_predictive_vars)
    354 ):
    355     priors_dict[group] = (
    356         None
    357         if var_names is None
--> 358         else dict_to_dataset(
    359             {k: np.expand_dims(self.prior[k], 0) for k in var_names},
    360             library=pymc,
    361             coords=self.coords,
    362             dims=self.dims,
    363         )
    364     )
    365 return priors_dict

File ...python/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 ...python/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 ...python/site-packages/arviz/data/base.py:254, in numpy_to_data_array(ary, var_name, coords, dims, default_dims, index_origin, skip_event_dims)
    251     coords["draw"] = np.arange(index_origin, n_samples + index_origin)
    253 # filter coords based on the dims
--> 254 coords = {key: xr.IndexVariable((key,), data=np.asarray(coords[key])) for key in dims}
    255 return xr.DataArray(ary, coords=coords, dims=dims)

File ...python/site-packages/arviz/data/base.py:254, in <dictcomp>(.0)
    251     coords["draw"] = np.arange(index_origin, n_samples + index_origin)
    253 # filter coords based on the dims
--> 254 coords = {key: xr.IndexVariable((key,), data=np.asarray(coords[key])) for key in dims}
    255 return xr.DataArray(ary, coords=coords, dims=dims)

File ...python/site-packages/xarray/core/variable.py:2846, in IndexVariable.__init__(self, dims, data, attrs, encoding, fastpath)
   2845 def __init__(self, dims, data, attrs=None, encoding=None, fastpath=False):
-> 2846     super().__init__(dims, data, attrs, encoding, fastpath)
   2847     if self.ndim != 1:
   2848         raise ValueError(f"{type(self).__name__} objects must be 1-dimensional")

File ...python/site-packages/xarray/core/variable.py:366, in Variable.__init__(self, dims, data, attrs, encoding, fastpath)
    346 """
    347 Parameters
    348 ----------
   (...)
    363     unrecognized encoding items.
    364 """
    365 self._data = as_compatible_data(data, fastpath=fastpath)
--> 366 self._dims = self._parse_dimensions(dims)
    367 self._attrs = None
    368 self._encoding = None

File ...python/site-packages/xarray/core/variable.py:663, in Variable._parse_dimensions(self, dims)
    661 dims = tuple(dims)
    662 if len(dims) != self.ndim:
--> 663     raise ValueError(
    664         f"dimensions {dims} must have the same length as the "
    665         f"number of data dimensions, ndim={self.ndim}"
    666     )
    667 return dims

ValueError: dimensions ('c',) must have the same length as the number of data dimensions, ndim=2

PyMC version information:

pymc : 5.0.2
arviz : 0.14.0
xarray: 2023.1.0

Context for the issue:

It looks like the error arises because along the way, the coord is converted from a tuple to a numpy array by InferenceDataConverter.

pymc/pymc/backends/arviz.py

Lines 210 to 214 in 902b1ec

	self.coords = {
	cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
	for cname, cvals in untyped_coords.items()
	if cvals is not None
	}

And numpy is eager to convert the coord vector into a multidimensional array, because it can. But np.array([("a", 1), ("a", 3), ("b", 2)]).ndim is 2, which I think causes the mismatch that makes arviz/xarray fail.

Obviously one workaround would be to stringify the coordinate values (that I currently have tupled). But it would be nice to preserve the type/objects.

[covertg]

If this is considered a bug -- could it be addressed by replacing np.array in inferencedataconverter with something like the following?

def array1d(x: Sequence):
    arr = np.array(x)
    if arr.ndim > 1:
        arr = np.empty(len(x), dtype="O")
        arr[:] = x
    return arr

This way numpy still infers the dtype unless the coord vals are sequences.

[michaelosthege]

I think a list(cvals) instead of np.array(cvals) should do the trick.
This would leave the responsibility for correct conversion to ArviZ.

watn to open a PR ?

[covertg]

OK, I've looked into a bit more. It's a bit hairy for a couple of reasons.

• xarray's IndexVariable (which backs each dim) ultimately calls asarray on any data it receives (in our case, pymc model coordinate values), which can lead to ndim > 1. https://github.com/pydata/xarray/blob/f46cd708f8e220272173e2fc3e66c7688df45c39/xarray/core/variable.py#L283
• Prior to this, arviz also calls asarray on coord data before passing it to IndexVariable https://github.com/arviz-devs/arviz/blob/631bbfedcea2ea6e1344ff5a490bc205aa1b8484/arviz/data/base.py#L254

So I think that each of those instances would have to be changed before list(cvals) in pymc will make any difference. I could get those conversations started upstream and see where it takes us.

Meanwhile, if the coord data is coerced into a 1dim numpy array anywhere along the call stack, it will remain that way. Thoughts on doing that as a fix on the pymc side or now? (So array1d(cvals) over list(cvals))

[covertg]

Here's a more direct way of playing with this:

import numpy as np
import pymc as pm
model = pm.Model(coords={"c": [("a", 1), ("a", 2), ("b", 1)]})
idc = pm.backends.arviz.InferenceDataConverter(
    model=model
    prior={"x": np.zeros((100, 3))}  # draws, ncoordvals,
    dims={"x": ["c"]}
)

idc.coords
# "{'c': array([['a', '1'], ['a', '2'], ['b', '1']], dtype='<U21')}"
idc.priors_to_xarray()
# same stack trace as above
idc.coords = {k: list(cval) for k, cval in model.coords.items()}
idc.priors_to_xarray()
# same

def array1d(x):
    arr = np.array(x)
    if arr.ndim > 1:
        arr = np.empty(len(x), dtype="O")
        arr[:] = x
    return arr
idc.coords = {k: array1d(cval) for k, cval in model.coords.items()}
idc.priors_to_xarray()
# works

[michaelosthege]

I think this is a similar situation: #6504 (comment)

From my experience the only true solution is to stop relying on NumPy magically knowing what's the right call.
For example, here I'm doing a lot of extra work to make sure that the ndarrays returned by this method are always of the expected shape and dtype based on metadata that is available ahead of time.
For example, if the shape metadata indicates that the items can have non-identical shapes the resulting ndarray should be dtype=object even if it's empty, len-1 or contains items with coincidentally identical shapes.

For the coordinate values problem here, this would mean that at the earliest possibility we should determine the appropriate dtype+shape and stick with it.
I the call stack that would be somewhere near Model.add_coord, and the Model.coords property could convert that to a Dict[str, tuple] to maintain the current signature.

[covertg]

Ahh, I see. Thanks for the references.

Updating add_coord and the coord property and some connected parts of Model seems not too hard. Thinking about a PR on this, though, I've noticed that the majority of references to Model.coords throughout the library end up converting the coordinate values to numpy... so it's relying on numpy magically knowing the right call in many places.

Suppose Model._coords is updated to be backed by a numpy array where we've been careful about shape and dtype. If other parts of the codebase want numpy for coordinate values, should they directly access (privately-named) _coords?

Other directions I can think of:

• Expose a new property in Model such as Model.coords_typed with signature Dict[str, np.ndarray]. Or Dict[str, pd.Index], see pm.Model(coords=coords) should preserve coordinate type if it is a pandas Index #6081
• Update the existing coords property to preserve shape and dtype as above, and update any code that relied on the tuple signature to tupleify the coord values, {cname: tuple(cvals) for cname, cvals in model.coords.items()}.
• Prevent the user from even using coordinate values that could be magically converted into a >1-dim array. E.g. update add_coord to throw an error if np.shape(cval) isn't length 1.

pymc/pymc/backends/arviz.py

Lines 219 to 227 in c5e4497

	# 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
	}

pymc/pymc/distributions/bound.py

Lines 197 to 199 in c5e4497

	if dims in model.coords:
	dim_obj = np.asarray(model.coords[dims])
	size = dim_obj.shape

pymc/pymc/model.py

Line 1011 in c5e4497

shape.extend(np.shape(self.coords[dim]))

pymc/pymc/sampling/jax.py

Lines 383 to 387 in c5e4497

	coords = {
	cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
	for cname, cvals in model.coords.items()
	if cvals is not None
	}

pymc/pymc/stats/log_likelihood.py

Lines 116 to 126 in c5e4497

	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()
	},
	default_dims=list(sample_dims),
	skip_event_dims=True,
	)

pymc/pymc/variational/opvi.py

Line 1134 in c5e4497

coords = {d: np.array(self.model.coords[d]) for d in dims}

[michaelosthege]

Like I explained in #6081 I don't like the idea of working with pd.Index objects because it is not clear how they would be stored without relying on xarray. If someone can propose how to modify this Coordinate proto to store pd.Index objects I'd be interested.

The current API of coords: Dict[str, Tuple] has the advantage that users can't edit the items. (Side note return self._coords should be changed to return {**self._coords}...)
If we decide to switch to Dict[str, np.ndarray] or any other type, we should make sure to return copies.