Implement flexible shape/dims/size API

RELEASE-NOTES.md

@@ -8,11 +8,18 @@
 
 ### New Features
 - The `CAR` distribution has been added to allow for use of conditional autoregressions which often are used in spatial and network models.
+- The dimensionality of model variables can now be parametrized through either of `shape`, `dims` or `size` (see [#4625](https://github.com/pymc-devs/pymc3/pull/4625)):
+  - With `shape` the length of dimensions must be given numerically or as scalar Aesara `Variables`. A `SpecifyShape` `Op` is added automatically unless `Ellipsis` is used. Using `shape` restricts the model variable to the exact length and re-sizing is no longer possible.
+  - `dims` keeps model variables re-sizeable (for example through `pm.Data`) and leads to well defined coordinates in `InferenceData` objects.
+  - The `size` kwarg creates new dimensions in addition to what is implied by RV parameters.
+  - An `Ellipsis` (`...`) in the last position of `shape` or `dims` can be used as short-hand notation for implied dimensions.
 - ...
 
 ### Maintenance
 - Remove float128 dtype support (see [#4514](https://github.com/pymc-devs/pymc3/pull/4514)).
 - Logp method of `Uniform` and `DiscreteUniform` no longer depends on `pymc3.distributions.dist_math.bound` for proper evaluation (see [#4541](https://github.com/pymc-devs/pymc3/pull/4541)).
+- `Model.RV_dims` and `Model.coords` are now read-only properties. To modify the `coords` dictionary use `Model.add_coord`. Also `dims` or coordinate values that are `None` will be auto-completed (see [#4625](https://github.com/pymc-devs/pymc3/pull/4625)).
+- The length of `dims` in the model is now tracked symbolically through `Model.dim_lengths` (see [#4625](https://github.com/pymc-devs/pymc3/pull/4625)).
 - ...
 
 ## PyMC3 3.11.2 (14 March 2021)

pymc3/aesaraf.py

@@ -45,6 +45,7 @@
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
+from aesara.tensor.shape import SpecifyShape
 from aesara.tensor.sharedvar import SharedVariable
 from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
 from aesara.tensor.var import TensorVariable
@@ -146,6 +147,8 @@ def change_rv_size(
         Expand the existing size by `new_size`.
 
     """
+    if isinstance(rv_var.owner.op, SpecifyShape):
+        rv_var = rv_var.owner.inputs[0]
     rv_node = rv_var.owner
     rng, size, dtype, *dist_params = rv_node.inputs
     name = rv_var.name

pymc3/backends/arviz.py

@@ -162,10 +162,7 @@ def __init__(
         self.trace = trace
 
         # this permits us to get the model from command-line argument or from with model:
-        try:
-            self.model = modelcontext(model)
-        except TypeError:
-            self.model = None
+        self.model = modelcontext(model)
 
         self.attrs = None
         if trace is not None:
@@ -223,10 +220,14 @@ def arbitrary_element(dct: Dict[Any, np.ndarray]) -> np.ndarray:
         self.coords = {} if coords is None else coords
         if hasattr(self.model, "coords"):
             self.coords = {**self.model.coords, **self.coords}
+        self.coords = {key: value for key, value in self.coords.items() if value is not None}
 
         self.dims = {} if dims is None else dims
         if hasattr(self.model, "RV_dims"):
-            model_dims = {k: list(v) for k, v in self.model.RV_dims.items()}
+            model_dims = {
+                var_name: [dim for dim in dims if dim is not None]
+                for var_name, dims in self.model.RV_dims.items()
+            }
             self.dims = {**model_dims, **self.dims}
 
         self.density_dist_obs = density_dist_obs

pymc3/data.py

@@ -19,7 +19,7 @@
 import urllib.request
 
 from copy import copy
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Sequence
 
 import aesara
 import aesara.tensor as at
@@ -502,7 +502,7 @@ class Data:
     >>> for data_vals in observed_data:
     ...     with model:
     ...         # Switch out the observed dataset
-    ...         pm.set_data({'data': data_vals})
+    ...         model.set_data('data', data_vals)
     ...         traces.append(pm.sample())
 
     To set the value of the data container variable, check out
@@ -543,6 +543,11 @@ def __new__(self, name, value, *, dims=None, export_index_as_coords=False):
 
         if export_index_as_coords:
             model.add_coords(coords)
+        elif dims:
+            # Register new dimension lengths
+            for d, dname in enumerate(dims):
+                if not dname in model.dim_lengths:
+                    model.add_coord(dname, values=None, length=shared_object.shape[d])
 
         # To draw the node for this variable in the graphviz Digraph we need
         # its shape.
@@ -562,7 +567,7 @@ def __new__(self, name, value, *, dims=None, export_index_as_coords=False):
         return shared_object
 
     @staticmethod
-    def set_coords(model, value, dims=None):
+    def set_coords(model, value, dims=None) -> Dict[str, Sequence]:
         coords = {}
 
         # If value is a df or a series, we interpret the index as coords: