Fix coords in constant_data issue #5046

[michaelosthege]

I got a minimum example to reproduce the problem.

@OriolAbril it looks like ArviZ 0.11.4 also has a bug there. This is the traceback when passing to dict_to_dataset from ArviZ:

_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _  

shape = (1, 17)
coords = {'backwards': <xarray.IndexVariable 'backwards' (backwards: 17)>
array([16, 15, 14, 13, 12, 11, 10,  9,  8,  7,  6,  5...: <xarray.IndexVariable 'draw' (draw: 17)>
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16])}
dims = ['chain', 'draw', 'backwards']

    def _infer_coords_and_dims(
        shape, coords, dims
    ) -> "Tuple[Dict[Any, Variable], Tuple[Hashable, ...]]":
        """All the logic for creating a new DataArray"""

        if (
            coords is not None
            and not utils.is_dict_like(coords)
            and len(coords) != len(shape)
        ):
            raise ValueError(
                f"coords is not dict-like, but it has {len(coords)} items, "
                f"which does not match the {len(shape)} dimensions of the "
                "data"
            )

        if isinstance(dims, str):
            dims = (dims,)

        if dims is None:
            dims = [f"dim_{n}" for n in range(len(shape))]
            if coords is not None and len(coords) == len(shape):
                # try to infer dimensions from coords
                if utils.is_dict_like(coords):
                    # deprecated in GH993, removed in GH1539
                    raise ValueError(
                        "inferring DataArray dimensions from "
                        "dictionary like ``coords`` is no longer "
                        "supported. Use an explicit list of "
                        "``dims`` instead."
                    )
                for n, (dim, coord) in enumerate(zip(dims, coords)):
                    coord = as_variable(coord, name=dims[n]).to_index_variable()
                    dims[n] = coord.name
            dims = tuple(dims)
        elif len(dims) != len(shape):
>           raise ValueError(
                "different number of dimensions on data "
                f"and dims: {len(shape)} vs {len(dims)}"
            )
E           ValueError: different number of dimensions on data and dims: 2 vs 3

..\..\appdata\local\continuum\miniconda3\envs\pm3v4\lib\site-packages\xarray\core\dataarray.py:126: ValueError
================================================================================================== warnings summary ================================================================================================== 
..\..\appdata\local\continuum\miniconda3\envs\pm3v4\lib\site-packages\win32\lib\pywintypes.py:2
  c:\users\osthege\appdata\local\continuum\miniconda3\envs\pm3v4\lib\site-packages\win32\lib\pywintypes.py:2: DeprecationWarning: the imp module is deprecated in favour of importlib; see the module's documentation for alternative uses
    import imp, sys, os

pymc/tests/test_idata_conversion.py::TestDataPyMC::test_constant_data_coords_issue_5046
  c:\users\osthege\appdata\local\continuum\miniconda3\envs\pm3v4\lib\site-packages\arviz\data\base.py:130: UserWarning: In variable alpha, there are more dims (1) given than exist (0). Passed array should have shape (chain,draw, *shape)
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/warnings.html
============================================================================================== short test summary info =============================================================================================== 
FAILED pymc/tests/test_idata_conversion.py::TestDataPyMC::test_constant_data_coords_issue_5046 - ValueError: different number of dimensions on data and dims: 2 vs 3

[codecov]

Codecov Report

Merging #5062 (b3daa2d) into main (a3cc81c) will decrease coverage by 20.36%.
The diff coverage is 25.00%.

@@             Coverage Diff             @@
##             main    #5062       +/-   ##
===========================================
- Coverage   78.32%   57.96%   -20.37%     
===========================================
  Files         130      130               
  Lines       24430    24439        +9     
===========================================
- Hits        19135    14165     -4970     
- Misses       5295    10274     +4979     

Impacted Files	Coverage Δ
pymc/tests/test_idata_conversion.py	0.00% <0.00%> (-96.21%)	⬇️
pymc/backends/arviz.py	75.09% <100.00%> (-14.34%)	⬇️
pymc/tests/test_ode.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/test_profile.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/test_model_func.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/test_posdef_sym.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/test_model_graph.py	0.00% <0.00%> (-100.00%)	⬇️
pymc/tests/test_model.py	0.00% <0.00%> (-99.12%)	⬇️
pymc/tests/test_shape_handling.py	0.00% <0.00%> (-96.03%)	⬇️
... and 30 more

[michaelosthege]

This bug was a nightmare. The workarounds in my plotting code were fragile and the debugging took ages. Too much staring at the bug without seeing it.

@OriolAbril @twiecki it would be awesome if we can fast-track this fix - it's blocking me quite severely.

[michaelosthege]

This red test is a flaky NUTS linear algebra tests (#5063). It has nothing to do with the changes from this PR, so I'm going to merge anyway.