Dataset.dims -> Dataset.sizes

Author: ricardoV94

pymc_experimental/tests/test_blackjax_smc.py

@@ -181,7 +181,7 @@ def test_arviz_from_particles():
     with model:
         inference_data = arviz_from_particles(model, particles)
 
-    assert inference_data.posterior.dims == Frozen({"chain": 1, "draw": 3, "x_dim_0": 2})
+    assert inference_data.posterior.sizes == Frozen({"chain": 1, "draw": 3, "x_dim_0": 2})
     assert inference_data.posterior.data_vars.dtypes == Frozen(
         {"x": dtype("float64"), "z": dtype("float64")}
     )

pymc_experimental/tests/test_linearmodel.py

@@ -142,8 +142,8 @@ def test_predict_posterior(fitted_linear_model_instance, combined):
     n_pred = 150
     X_pred = pd.DataFrame({"input": np.random.uniform(low=0, high=1, size=n_pred)})
     pred = model.predict_posterior(X_pred, combined=combined)
-    chains = model.idata.sample_stats.dims["chain"]
-    draws = model.idata.sample_stats.dims["draw"]
+    chains = model.idata.sample_stats.sizes["chain"]
+    draws = model.idata.sample_stats.sizes["draw"]
     expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
     assert pred.shape == expected_shape
     assert np.issubdtype(pred.dtype, np.floating)

pymc_experimental/tests/test_model_builder.py

@@ -238,8 +238,8 @@ def test_sample_posterior_predictive(fitted_model_instance, combined):
     pred = fitted_model_instance.sample_posterior_predictive(
         prediction_data["input"], combined=combined, extend_idata=True
     )
-    chains = fitted_model_instance.idata.sample_stats.dims["chain"]
-    draws = fitted_model_instance.idata.sample_stats.dims["draw"]
+    chains = fitted_model_instance.idata.sample_stats.sizes["chain"]
+    draws = fitted_model_instance.idata.sample_stats.sizes["draw"]
     expected_shape = (n_pred, chains * draws) if combined else (chains, draws, n_pred)
     assert pred[fitted_model_instance.output_var].shape == expected_shape
     assert np.issubdtype(pred[fitted_model_instance.output_var].dtype, np.floating)