Check idata dims/coords for consistency before sampling begins

[jessegrabowski]

Description

This is a painful gotcha I get hit by once in a while. If you have inconsistencies between the model coords and variable shapes, PyMC will still happy sample (often for hours, ask me how I know), then error out when it gets to the final step of building the idata.

I don't think there needs to be complex logic to validate static shapes etc, I'm thinking we could do a smoke test on the model initial point (or a single prior draw) to make sure it can be converted to an idata, and raise before sampling if not.

[ricardoV94]

I'm against this. It's really nice to have sampling responsive, I had to cut so much cruft in #7578, including a check_shapes call for every chain/ip.

The smoke test should be by the user, just sample a few tune/draws if you're unsure? Or have arviz fail gracefully (which I think we do now ?)

[ricardoV94]

Don't you trigger this in your case?

pymc/pymc/backends/arviz.py

Lines 60 to 86 in 0960323

	def dict_to_dataset_drop_incompatible_coords(vars_dict, *args, dims, coords, **kwargs):
	safe_coords = coords
	if not RAISE_ON_INCOMPATIBLE_COORD_LENGTHS:
	coords_lengths = {k: len(v) for k, v in coords.items()}
	for var_name, var in vars_dict.items():
	# Iterate in reversed because of chain/draw batch dimensions
	for dim, dim_length in zip(reversed(dims.get(var_name, ())), reversed(var.shape)):
	coord_length = coords_lengths.get(dim, None)
	if (coord_length is not None) and (coord_length != dim_length):
	warnings.warn(
	f"Incompatible coordinate length of {coord_length} for dimension '{dim}' of variable '{var_name}'.\n"
	"This usually happens when a sliced or concatenated variable is wrapped as a `pymc.dims.Deterministic`."
	"The originate coordinates for this dim will not be included in the returned dataset for any of the variables. "
	"Instead they will default to `np.arange(var_length)` and the shorter variables will be right-padded with nan.\n"
	"To make this warning into an error set `pymc.backends.arviz.RAISE_ON_INCOMPATIBLE_COORD_LENGTHS` to `True`",
	UserWarning,
	)
	if safe_coords is coords:
	safe_coords = coords.copy()
	safe_coords.pop(dim)
	coords_lengths.pop(dim)
	# FIXME: Would be better to drop coordinates altogether, but arviz defaults to `np.arange(var_length)`
	return dict_to_dataset(vars_dict, *args, dims=dims, coords=safe_coords, **kwargs)

[jessegrabowski]

Let me cook up an MWE. If we can gracefully fail to convert to arviz and return a raw trace object that would also be good. I just think it's too punishing to do a long sampling run, have it succeed, and lose all outputs because of a stupid error.

[ricardoV94]

Agreed