Option to create non-shared pm.Data

RELEASE-NOTES.md

@@ -113,7 +113,10 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
   - Added partial dependence plots and individual conditional expectation plots [5091](https://github.com/pymc-devs/pymc3/pull/5091).
   - Modify how particle weights are computed. This improves accuracy of the modeled function (see [5177](https://github.com/pymc-devs/pymc3/pull/5177)).
   - Improve sampling, increase default number of particles [5229](https://github.com/pymc-devs/pymc3/pull/5229).
-- `pm.Data` now passes additional kwargs to `aesara.shared`. [#5098](https://github.com/pymc-devs/pymc/pull/5098)
+- New features for `pm.Data` containers:
+  - With `pm.Data(..., mutable=True/False)`, or by using `pm.MutableData` vs. `pm.ConstantData` one can now create `TensorConstant` data variables. They can be more performant and compatible in situtations where a variable doesn't need to be changed via `pm.set_data()`. See [#5295](https://github.com/pymc-devs/pymc/pull/5295).
+  - New named dimensions can be introduced to the model via `pm.Data(..., dims=...)`. For mutable data variables (see above) the lengths of these dimensions are symbolic, so they can be re-sized via `pm.set_data()`.
+  - `pm.Data` now passes additional kwargs to `aesara.shared`/`at.as_tensor`. [#5098](https://github.com/pymc-devs/pymc/pull/5098).
 - ...
 
 

pymc/data.py

@@ -19,16 +19,17 @@
 import urllib.request
 
 from copy import copy
-from typing import Any, Dict, List, Sequence
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import aesara
 import aesara.tensor as at
 import numpy as np
 import pandas as pd
 
+from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph.basic import Apply
 from aesara.tensor.type import TensorType
-from aesara.tensor.var import TensorVariable
+from aesara.tensor.var import TensorConstant, TensorVariable
 
 import pymc as pm
 
@@ -40,6 +41,8 @@
     "Minibatch",
     "align_minibatches",
     "Data",
+    "ConstantData",
+    "MutableData",
 ]
 BASE_URL = "https://github.com/raw/pymc-devs/pymc-examples/main/examples/data/{filename}"
 
@@ -463,9 +466,103 @@ def align_minibatches(batches=None):
                 rng.seed()
 
 
-class Data:
-    """Data container class that wraps :func:`aesara.shared` and lets
-    the model be aware of its inputs and outputs.
+def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict[str, Sequence]:
+    """Determines coordinate values from data or the model (via ``dims``)."""
+    coords = {}
+
+    # If value is a df or a series, we interpret the index as coords:
+    if isinstance(value, (pd.Series, pd.DataFrame)):
+        dim_name = None
+        if dims is not None:
+            dim_name = dims[0]
+        if dim_name is None and value.index.name is not None:
+            dim_name = value.index.name
+        if dim_name is not None:
+            coords[dim_name] = value.index
+
+    # If value is a df, we also interpret the columns as coords:
+    if isinstance(value, pd.DataFrame):
+        dim_name = None
+        if dims is not None:
+            dim_name = dims[1]
+        if dim_name is None and value.columns.name is not None:
+            dim_name = value.columns.name
+        if dim_name is not None:
+            coords[dim_name] = value.columns
+
+    if isinstance(value, np.ndarray) and dims is not None:
+        if len(dims) != value.ndim:
+            raise pm.exceptions.ShapeError(
+                "Invalid data shape. The rank of the dataset must match the " "length of `dims`.",
+                actual=value.shape,
+                expected=value.ndim,
+            )
+        for size, dim in zip(value.shape, dims):
+            coord = model.coords.get(dim, None)
+            if coord is None:
+                coords[dim] = pd.RangeIndex(size, name=dim)
+
+    return coords
+
+
+def ConstantData(
+    name: str,
+    value,
+    *,
+    dims: Optional[Sequence[str]] = None,
+    export_index_as_coords=False,
+    **kwargs,
+) -> TensorConstant:
+    """Alias for ``pm.Data(..., mutable=False)``.
+
+    Registers the ``value`` as a ``TensorConstant`` with the model.
+    """
+    return Data(
+        name,
+        value,
+        dims=dims,
+        export_index_as_coords=export_index_as_coords,
+        mutable=False,
+        **kwargs,
+    )
+
+
+def MutableData(
+    name: str,
+    value,
+    *,
+    dims: Optional[Sequence[str]] = None,
+    export_index_as_coords=False,
+    **kwargs,
+) -> SharedVariable:
+    """Alias for ``pm.Data(..., mutable=True)``.
+
+    Registers the ``value`` as a ``SharedVariable`` with the model.
+    """
+    return Data(
+        name,
+        value,
+        dims=dims,
+        export_index_as_coords=export_index_as_coords,
+        mutable=True,
+        **kwargs,
+    )
+
+
+def Data(
+    name: str,
+    value,
+    *,
+    dims: Optional[Sequence[str]] = None,
+    export_index_as_coords=False,
+    mutable: bool = True,
+    **kwargs,
+) -> Union[SharedVariable, TensorConstant]:
+    """Data container that registers a data variable with the model.
+
+    Depending on the ``mutable`` setting (default: True), the variable
+    is registered as a ``SharedVariable``, enabling it to be altered
+    in value and shape, but NOT in dimensionality using ``pm.set_data()``.
 
     Parameters
     ----------
@@ -513,104 +610,46 @@ class Data:
     For more information, take a look at this example notebook
     https://docs.pymc.io/notebooks/data_container.html
     """
+    if isinstance(value, list):
+        value = np.array(value)
 
-    def __new__(
-        self,
-        name,
-        value,
-        *,
-        dims=None,
-        export_index_as_coords=False,
-        **kwargs,
-    ):
-        if isinstance(value, list):
-            value = np.array(value)
-
-        # Add data container to the named variables of the model.
-        try:
-            model = pm.Model.get_context()
-        except TypeError:
-            raise TypeError(
-                "No model on context stack, which is needed to instantiate a data container. "
-                "Add variable inside a 'with model:' block."
-            )
-        name = model.name_for(name)
-
-        # `pandas_to_array` takes care of parameter `value` and
-        # transforms it to something digestible for pymc
-        shared_object = aesara.shared(pandas_to_array(value), name, **kwargs)
-
-        if isinstance(dims, str):
-            dims = (dims,)
-        if not (dims is None or len(dims) == shared_object.ndim):
-            raise pm.exceptions.ShapeError(
-                "Length of `dims` must match the dimensions of the dataset.",
-                actual=len(dims),
-                expected=shared_object.ndim,
-            )
-
-        coords = self.set_coords(model, value, dims)
-
-        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.
-        # XXX: This needs to be refactored
-        # shared_object.dshape = tuple(shared_object.shape.eval())
-        # if dims is not None:
-        #     shape_dims = model.shape_from_dims(dims)
-        #     if shared_object.dshape != shape_dims:
-        #         raise pm.exceptions.ShapeError(
-        #             "Data shape does not match with specified `dims`.",
-        #             actual=shared_object.dshape,
-        #             expected=shape_dims,
-        #         )
-
-        model.add_random_variable(shared_object, dims=dims)
-
-        return shared_object
-
-    @staticmethod
-    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:
-        if isinstance(value, (pd.Series, pd.DataFrame)):
-            dim_name = None
-            if dims is not None:
-                dim_name = dims[0]
-            if dim_name is None and value.index.name is not None:
-                dim_name = value.index.name
-            if dim_name is not None:
-                coords[dim_name] = value.index
-
-        # If value is a df, we also interpret the columns as coords:
-        if isinstance(value, pd.DataFrame):
-            dim_name = None
-            if dims is not None:
-                dim_name = dims[1]
-            if dim_name is None and value.columns.name is not None:
-                dim_name = value.columns.name
-            if dim_name is not None:
-                coords[dim_name] = value.columns
-
-        if isinstance(value, np.ndarray) and dims is not None:
-            if len(dims) != value.ndim:
-                raise pm.exceptions.ShapeError(
-                    "Invalid data shape. The rank of the dataset must match the "
-                    "length of `dims`.",
-                    actual=value.shape,
-                    expected=value.ndim,
-                )
-            for size, dim in zip(value.shape, dims):
-                coord = model.coords.get(dim, None)
-                if coord is None:
-                    coords[dim] = pd.RangeIndex(size, name=dim)
-
-        return coords
+    # Add data container to the named variables of the model.
+    try:
+        model = pm.Model.get_context()
+    except TypeError:
+        raise TypeError(
+            "No model on context stack, which is needed to instantiate a data container. "
+            "Add variable inside a 'with model:' block."
+        )
+    name = model.name_for(name)
+
+    # `pandas_to_array` takes care of parameter `value` and
+    # transforms it to something digestible for Aesara.
+    arr = pandas_to_array(value)
+    if mutable:
+        x = aesara.shared(arr, name, **kwargs)
+    else:
+        x = at.as_tensor_variable(arr, name, **kwargs)
+
+    if isinstance(dims, str):
+        dims = (dims,)
+    if not (dims is None or len(dims) == x.ndim):
+        raise pm.exceptions.ShapeError(
+            "Length of `dims` must match the dimensions of the dataset.",
+            actual=len(dims),
+            expected=x.ndim,
+        )
+
+    coords = determine_coords(model, value, dims)
+
+    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=x.shape[d])
+
+    model.add_random_variable(x, dims=dims)
+
+    return x

pymc/model_graph.py

@@ -20,7 +20,7 @@
 from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph.basic import walk
 from aesara.tensor.random.op import RandomVariable
-from aesara.tensor.var import TensorVariable
+from aesara.tensor.var import TensorConstant, TensorVariable
 
 import pymc as pm
 
@@ -133,7 +133,7 @@ def _make_node(self, var_name, graph, *, formatting: str = "plain"):
             shape = "octagon"
             style = "filled"
             label = f"{var_name}\n~\nPotential"
-        elif isinstance(v, SharedVariable):
+        elif isinstance(v, (SharedVariable, TensorConstant)):
             shape = "box"
             style = "rounded, filled"
             label = f"{var_name}\n~\nData"

pymc/tests/test_data_container.py

@@ -71,8 +71,8 @@ def test_sample(self):
 
     def test_sample_posterior_predictive_after_set_data(self):
         with pm.Model() as model:
-            x = pm.Data("x", [1.0, 2.0, 3.0])
-            y = pm.Data("y", [1.0, 2.0, 3.0])
+            x = pm.MutableData("x", [1.0, 2.0, 3.0])
+            y = pm.ConstantData("y", [1.0, 2.0, 3.0])
             beta = pm.Normal("beta", 0, 10.0)
             pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y)
             trace = pm.sample(