ModelReviewer/BayesianPPPCheck fails when using JAX backend - `TypeError: einsum requires ndarray or scalar arguments, got xarray at position 0.`

[aswan89]

When executing meridian.analysis.review.reviewer.ModelReviewer(mmm).run() using the JAX backend end I receive the error ypeError: einsum requires ndarray or scalar arguments, got xarray at position 0.

Going through the code I believe the issue is here, line 181 of checks.py. model_context.input_data.revenue_per_kpi is a xarray object and multiplying kpi it coerces kpi to Xarray, but analyzer.filter_and_aggregate_geos_and_times is expecting backend.Tensor. I can't say for certain that the tensorflow backend can coerce xarray to something einsum can handle but it seems that JAX cannot.

For now I'm doing a workaround where I manually cast revenue_per_kpi via mmm.model_context.input_data.revenue_per_kpi = mmm.model_context.input_data.revenue_per_kpi.to_numpy()

[cpulavarthi]

Hello @aswan89,

Thank you for bringing this issue to our notice. Our team has confirmed the bug and is working towards resolving it. This issue should be fixed in the next release of Meridian. You can track the previous changes to Meridian and unreleased feature updates in our Change Log.

Feel free to reach out if you have any questions or suggestions regarding Meridian.

Google Meridian Support Team