.more stuff

Author: ricardoV94

docs/source/learn/core_notebooks/dims_module.ipynb

@@ -49,9 +49,37 @@
 
 _log = logging.getLogger(__name__)
 
+
+RAISE_ON_INCOMPATIBLE_COORD_LENGTHS = False
+
+
 # random variable object ...
 Var = Any
 
+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[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} found for dimension {dim} of variable {var_name}.\n"
+                        "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)`.\n"
+                        "To make this warning into an errror 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)
 
 def find_observations(model: "Model") -> dict[str, Var]:
     """If there are observations available, return them as a dictionary."""
@@ -365,7 +393,7 @@ def priors_to_xarray(self):
             priors_dict[group] = (
                 None
                 if var_names is None
-                else dict_to_dataset(
+                else dict_to_dataset_drop_incompatible_coords(
                     {k: np.expand_dims(self.prior[k], 0) for k in var_names},
                     library=pymc,
                     coords=self.coords,

pymc/backends/arviz.py

@@ -68,10 +68,13 @@ def extend_dim(array, dim):
         fill_val = norm - sum_vals / pt.sqrt(n)
 
         out = ptx.concat([array, fill_val], dim=dim)
+        import pytensor
+        with pytensor.config.change_flags(optimizer_verbose=True):
+            print(out.eval(mode="minimum_compile").shape)
         return out - norm
 
     @staticmethod
-    def reduct_dim(array, dim):
+    def reduce_dim(array, dim):
         n = array.sizes[dim].astype("floatX")
         last = array.isel({dim: -1})
 
@@ -81,7 +84,7 @@ def reduct_dim(array, dim):
 
     def forward(self, value, *rv_inputs):
         for dim in self.dims:
-            value = self.reduct_dim(value, dim=dim)
+            value = self.reduce_dim(value, dim=dim)
         return value
 
     def backward(self, value, *rv_inputs):
@@ -92,4 +95,4 @@ def backward(self, value, *rv_inputs):
     def log_jac_det(self, value, *rv_inputs):
         # Use following once broadcast_like is implemented
         # as_xtensor(0).broadcast_like(value, exclude=self.dims)`
-        return (value * 0).sum(self.dims)
+        return value.sum(self.dims) * 0

pymc/dims/transforms.py

@@ -1781,7 +1781,7 @@ def point_logps(self, point=None, round_vals=2, **kwargs):
             point = self.initial_point()
 
         factors = self.basic_RVs + self.potentials
-        factor_logps_fn = [pt.sum(factor) for factor in self.logp(factors, sum=False)]
+        factor_logps_fn = [factor.sum() for factor in self.logp(factors, sum=False)]
         return {
             factor.name: np.round(np.asarray(factor_logp), round_vals)
             for factor, factor_logp in zip(

pymc/model/core.py

@@ -77,8 +77,11 @@ def create_plate_label_with_dim_length(
     )
 
 
+from pymc.pytensorf import _cheap_eval_mode
+
+
 def fast_eval(var):
-    return function([], var, mode="FAST_COMPILE")()
+    return function([], var, mode=_cheap_eval_mode)()
 
 
 class NodeType(str, Enum):

pymc/model_graph.py

@@ -60,7 +60,9 @@ def test_zerosumnormal():
         np.testing.assert_allclose(ip_value, ref_ip_value)
 
     logp_fn = model.compile_logp()
-    ref_logp_fn = ref_model.compile_logp()
     logp_fn(ip)
+    ref_logp_fn = ref_model.compile_logp()
     # np.testing.assert_allclose(logp_fn(ip), ref_logp_fn(ref_ip))
     # Test a new
+
+test_zerosumnormal()