Update Aesara to 2.8.8

This includes several changes related to static shape handling.

`test_var_replacement` is marked XFAIL because of incompatibility
with the new static shape handling.

Author: michaelosthege

conda-envs/environment-dev.yml

@@ -5,7 +5,7 @@ channels:
 - defaults
 dependencies:
 # Base dependencies
-- aesara=2.8.7
+- aesara=2.8.8
 - arviz>=0.13.0
 - blas
 - cachetools>=4.2.1

conda-envs/environment-test.yml

@@ -5,7 +5,7 @@ channels:
 - defaults
 dependencies:
 # Base dependencies
-- aesara=2.8.7
+- aesara=2.8.8
 - arviz>=0.13.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-dev.yml

@@ -5,7 +5,7 @@ channels:
 - defaults
 dependencies:
 # Base dependencies (see install guide for Windows)
-- aesara=2.8.7
+- aesara=2.8.8
 - arviz>=0.13.0
 - blas
 - cachetools>=4.2.1

conda-envs/windows-environment-test.yml

@@ -5,7 +5,7 @@ channels:
 - defaults
 dependencies:
 # Base dependencies (see install guide for Windows)
-- aesara=2.8.7
+- aesara=2.8.8
 - arviz>=0.13.0
 - blas
 - cachetools>=4.2.1

pymc/math.py

@@ -284,15 +284,15 @@ def softmax(x, axis=None):
     # drops that warning
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", UserWarning)
-        return at.nnet.softmax(x, axis=axis)
+        return at.special.softmax(x, axis=axis)
 
 
 def log_softmax(x, axis=None):
     # Ignore vector case UserWarning issued by Aesara. This can be removed once Aesara
     # drops that warning
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", UserWarning)
-        return at.nnet.logsoftmax(x, axis=axis)
+        return at.special.log_softmax(x, axis=axis)
 
 
 def logbern(log_p):

pymc/model.py

@@ -1212,15 +1212,29 @@ def set_data(
                         "or define it via a `pm.MutableData` variable."
                     )
                 elif length_tensor.owner is not None:
-                    # The dimension was created from a model variable.
+                    # The dimension was created from another variable:
+                    length_tensor_origin = length_tensor.owner.inputs[0]
                     # Get a handle on the tensor from which this dimension length was
                     # obtained by doing subindexing on the shape as in `.shape[i]`.
-                    # Needed to check if it was another shared variable.
+                    if isinstance(length_tensor_origin, TensorConstant):
+                        raise ShapeError(
+                            f"Resizing dimension '{dname}' with values of length {new_length} would lead to incompatibilities, "
+                            f"because the dimension length is tied to a {length_tensor_origin}. "
+                            f"Check if the dimension was defined implicitly before the shared variable '{name}' was created, "
+                            f"for example by another model variable.",
+                            actual=new_length,
+                            expected=old_length,
+                        )
+
+                    # The shape entry this dimension is tied to is not a TensorConstant.
+                    # Whether the dimension can be resized depends on the kind of Variable the shape belongs to.
                     # TODO: Consider checking the graph is what we are assuming it is
                     # isinstance(length_tensor.owner.op, Subtensor)
                     # isinstance(length_tensor.owner.inputs[0].owner.op, Shape)
-                    length_belongs_to = length_tensor.owner.inputs[0].owner.inputs[0]
+                    length_belongs_to = length_tensor_origin.owner.inputs[0]
+
                     if length_belongs_to is shared_object:
+                        # This is the shared variable that's being updated!
                         # No surprise it's changing.
                         pass
                     elif isinstance(length_belongs_to, SharedVariable):

pymc/tests/distributions/test_truncated.py

@@ -97,7 +97,7 @@ def test_truncation_continuous_random(op_type, lower, upper):
 
     xt = Truncated.dist(x, lower=lower, upper=upper)
     assert isinstance(xt.owner.op, TruncatedRV)
-    assert xt.type == x.type
+    assert xt.type.dtype == x.type.dtype
 
     xt_draws = draw(xt, draws=5)
     assert np.all(xt_draws >= lower)
@@ -162,7 +162,7 @@ def test_truncation_discrete_random(op_type, lower, upper):
     x = geometric_op(p, name="x", size=500)
     xt = Truncated.dist(x, lower=lower, upper=upper)
     assert isinstance(xt.owner.op, TruncatedRV)
-    assert xt.type == x.type
+    assert xt.type.dtype == x.type.dtype
 
     xt_draws = draw(xt)
     assert np.all(xt_draws >= lower)

pymc/tests/test_data.py

@@ -714,7 +714,6 @@ def test_free_rv(self):
         with pm.Model() as model5:
             n = pm.Normal("n", total_size=[2, Ellipsis, 2], size=(2, 2))
             p5 = model5.compile_fn(model5.logp(), point_fn=False)
-        assert p4() == p5(pm.floatX([[1]]))
         assert p4() == p5(pm.floatX([[1, 1], [1, 1]]))
 
 

pymc/tests/test_math.py

@@ -280,8 +280,8 @@ def test_invlogit_deprecation_warning():
 @pytest.mark.parametrize(
     "aesara_function, pymc_wrapper",
     [
-        (at.nnet.softmax, softmax),
-        (at.nnet.logsoftmax, log_softmax),
+        (at.special.softmax, softmax),
+        (at.special.log_softmax, log_softmax),
     ],
 )
 def test_softmax_logsoftmax_no_warnings(aesara_function, pymc_wrapper):

pymc/tests/test_model.py

@@ -216,8 +216,8 @@ def test_observed_type(self):
             x1 = pm.Normal("x1", observed=X_)
             x2 = pm.Normal("x2", observed=X)
 
-        assert x1.type == X.type
-        assert x2.type == X.type
+        assert x1.type.dtype == X.type.dtype
+        assert x2.type.dtype == X.type.dtype
 
 
 def test_duplicate_vars():
@@ -935,7 +935,7 @@ def test_set_data_constant_shape_error():
         pmodel.add_coord("weekday", length=x.shape[0])
         pm.MutableData("y", np.arange(7), dims="weekday")
 
-    msg = "because the dimension was initialized from 'x' which is not a shared variable"
+    msg = "because the dimension length is tied to a TensorConstant"
     with pytest.raises(ShapeError, match=msg):
         pmodel.set_data("y", np.arange(10))
 

pymc/tests/variational/test_inference.py

@@ -307,6 +307,7 @@ def test_remove_scan_op():
         buff.close()
 
 
+@pytest.mark.xfail(reason="Broke from static shape handling with Aesara 2.8.8")
 def test_var_replacement():
     X_mean = pm.floatX(np.linspace(0, 10, 10))
     y = pm.floatX(np.random.normal(X_mean * 4, 0.05))

pymc/variational/updates.py

@@ -1052,7 +1052,7 @@ def total_norm_constraint(tensor_vars, max_norm, epsilon=1e-7, return_norm=False
     >>> x = at.matrix()
     >>> y = at.ivector()
     >>> l_in = InputLayer((5, 10))
-    >>> l1 = DenseLayer(l_in, num_units=7, nonlinearity=at.nnet.softmax)
+    >>> l1 = DenseLayer(l_in, num_units=7, nonlinearity=at.special.softmax)
     >>> output = lasagne.layers.get_output(l1, x)
     >>> cost = at.mean(at.nnet.categorical_crossentropy(output, y))
     >>> all_params = lasagne.layers.get_all_params(l1)

requirements-dev.txt

@@ -1,7 +1,7 @@
 # This file is auto-generated by scripts/generate_pip_deps_from_conda.py, do not modify.
 # See that file for comments about the need/usage of each dependency.
 
-aesara==2.8.7
+aesara==2.8.8
 arviz>=0.13.0
 cachetools>=4.2.1
 cloudpickle

requirements.txt

@@ -1,4 +1,4 @@
-aesara==2.8.7
+aesara==2.8.8
 arviz>=0.13.0
 cachetools>=4.2.1
 cloudpickle