Repository: aesara-devs/aesara

Issue:
Error w/ deepcopy of pymc v4 model
## Generating a deep copy of pymc v4 model results in aesara error

**Please provide a minimal, self-contained, and reproducible example.**
```python
from copy import deepcopy

import numpy as np
import pymc as pm

RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)

# True parameter values
alpha, sigma = 1, 1
beta = [1, 2.5]

# Size of dataset
size = 100

# Predictor variable
X1 = np.random.randn(size)
X2 = np.random.randn(size) * 0.2

# Simulate outcome variable
Y = alpha + beta[0] * X1 + beta[1] * X2 + rng.normal(size=size) * sigma

basic_model = pm.Model()

with basic_model:

    # Priors for unknown model parameters
    alpha = pm.Normal("alpha", mu=0, sigma=10)
    beta = pm.Normal("beta", mu=0, sigma=10, shape=2)
    sigma = pm.HalfNormal("sigma", sigma=1)

    # Expected value of outcome
    mu = alpha + beta[0] * X1 + beta[1] * X2

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)

model_copy = deepcopy(basic_model)
```

**Please provide the full traceback of any errors.**
```python
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [9], in <cell line: 41>()
     38     # Likelihood (sampling distribution) of observations
     39     Y_obs = pm.Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)
---> 41 model_copy = deepcopy(basic_model)

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:270, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    268 if state is not None:
    269     if deep:
--> 270         state = deepcopy(state, memo)
    271     if hasattr(y, '__setstate__'):
    272         y.__setstate__(state)

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:230, in _deepcopy_dict(x, memo, deepcopy)
    228 memo[id(x)] = y
    229 for key, value in x.items():
--> 230     y[deepcopy(key, memo)] = deepcopy(value, memo)
    231 return y

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:230, in _deepcopy_dict(x, memo, deepcopy)
    228 memo[id(x)] = y
    229 for key, value in x.items():
--> 230     y[deepcopy(key, memo)] = deepcopy(value, memo)
    231 return y

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:270, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    268 if state is not None:
    269     if deep:
--> 270         state = deepcopy(state, memo)
    271     if hasattr(y, '__setstate__'):
    272         y.__setstate__(state)

    [... skipping similar frames: _deepcopy_dict at line 230 (1 times), deepcopy at line 146 (1 times)]

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:270, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    268 if state is not None:
    269     if deep:
--> 270         state = deepcopy(state, memo)
    271     if hasattr(y, '__setstate__'):
    272         y.__setstate__(state)

    [... skipping similar frames: deepcopy at line 146 (1 times)]

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:230, in _deepcopy_dict(x, memo, deepcopy)
    228 memo[id(x)] = y
    229 for key, value in x.items():
--> 230     y[deepcopy(key, memo)] = deepcopy(value, memo)
    231 return y

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:205, in _deepcopy_list(x, memo, deepcopy)
    203 append = y.append
    204 for a in x:
--> 205     append(deepcopy(a, memo))
    206 return y

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:270, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    268 if state is not None:
    269     if deep:
--> 270         state = deepcopy(state, memo)
    271     if hasattr(y, '__setstate__'):
    272         y.__setstate__(state)

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:230, in _deepcopy_dict(x, memo, deepcopy)
    228 memo[id(x)] = y
    229 for key, value in x.items():
--> 230     y[deepcopy(key, memo)] = deepcopy(value, memo)
    231 return y

File ~\AppData\Local\Programs\Python\Python39\lib\copy.py:153, in deepcopy(x, memo, _nil)
    151 copier = getattr(x, "__deepcopy__", None)
    152 if copier is not None:
--> 153     y = copier(memo)
    154 else:
    155     reductor = dispatch_table.get(cls)

File ~\.virtualenvs\pymc4-venv\lib\site-packages\aesara\link\basic.py:142, in Container.__deepcopy__(self, memo)
    132 r = type(self)(
    133     deepcopy(self.type, memo=memo),
    134     deepcopy(self.storage, memo=memo),
   (...)
    138     name=deepcopy(self.name, memo=memo),
    139 )
    140 # Work around NumPy deepcopy of ndarray with 0 dimension that
    141 # don't return an ndarray.
--> 142 if r.storage[0] is not None and not self.type.is_valid_value(r.storage[0]):
    143     assert not data_was_in_memo
    144     assert self.type.is_valid_value(self.storage[0])

TypeError: is_valid_value() missing 1 required positional argument: 'strict'
```

This used to work on pymc3 and theano-pymc


## Versions and main components

* Aesara version: 2.66
* Aesara config (`python -c "import aesara; print(aesara.config)"`)
* Python version: 3.9
* Operating system:  Windows
* How did you install Aesara: pip
* Pymc version: 4.0


Comments:

Patch:
diff --git a/aesara/graph/type.py b/aesara/graph/type.py
index f26d72664..f35a6d969 100644
--- a/aesara/graph/type.py
+++ b/aesara/graph/type.py
@@ -180,10 +180,10 @@ class Type(MetaObject):
 
         return None
 
-    def is_valid_value(self, data: D) -> bool:
+    def is_valid_value(self, data: D, strict: bool = True) -> bool:
         """Return ``True`` for any python object that would be a legal value for a `Variable` of this `Type`."""
         try:
-            self.filter(data, strict=True)
+            self.filter(data, strict=strict)
             return True
         except (TypeError, ValueError):
             return False
diff --git a/aesara/tensor/random/type.py b/aesara/tensor/random/type.py
index ff00669c4..68373b08c 100644
--- a/aesara/tensor/random/type.py
+++ b/aesara/tensor/random/type.py
@@ -1,9 +1,14 @@
+from typing import Generic, TypeVar
+
 import numpy as np
 
 import aesara
 from aesara.graph.type import Type
 
 
+T = TypeVar("T", np.random.RandomState, np.random.Generator)
+
+
 gen_states_keys = {
     "MT19937": (["state"], ["key", "pos"]),
     "PCG64": (["state", "has_uint32", "uinteger"], ["state", "inc"]),
@@ -18,22 +23,15 @@ gen_states_keys = {
 numpy_bit_gens = {0: "MT19937", 1: "PCG64", 2: "Philox", 3: "SFC64"}
 
 
-class RandomType(Type):
+class RandomType(Type, Generic[T]):
     r"""A Type wrapper for `numpy.random.Generator` and `numpy.random.RandomState`."""
 
-    @classmethod
-    def filter(cls, data, strict=False, allow_downcast=None):
-        if cls.is_valid_value(data, strict):
-            return data
-        else:
-            raise TypeError()
-
     @staticmethod
-    def may_share_memory(a, b):
+    def may_share_memory(a: T, b: T):
         return a._bit_generator is b._bit_generator
 
 
-class RandomStateType(RandomType):
+class RandomStateType(RandomType[np.random.RandomState]):
     r"""A Type wrapper for `numpy.random.RandomState`.
 
     The reason this exists (and `Generic` doesn't suffice) is that
@@ -49,28 +47,38 @@ class RandomStateType(RandomType):
     def __repr__(self):
         return "RandomStateType"
 
-    @staticmethod
-    def is_valid_value(a, strict):
-        if isinstance(a, np.random.RandomState):
-            return True
+    def filter(self, data, strict: bool = False, allow_downcast=None):
+        """
+        XXX: This doesn't convert `data` to the same type of underlying RNG type
+        as `self`.  It really only checks that `data` is of the appropriate type
+        to be a valid `RandomStateType`.
+
+        In other words, it serves as a `Type.is_valid_value` implementation,
+        but, because the default `Type.is_valid_value` depends on
+        `Type.filter`, we need to have it here to avoid surprising circular
+        dependencies in sub-classes.
+        """
+        if isinstance(data, np.random.RandomState):
+            return data
 
-        if not strict and isinstance(a, dict):
+        if not strict and isinstance(data, dict):
             gen_keys = ["bit_generator", "gauss", "has_gauss", "state"]
             state_keys = ["key", "pos"]
 
             for key in gen_keys:
-                if key not in a:
-                    return False
+                if key not in data:
+                    raise TypeError()
 
             for key in state_keys:
-                if key not in a["state"]:
-                    return False
+                if key not in data["state"]:
+                    raise TypeError()
 
-            state_key = a["state"]["key"]
+            state_key = data["state"]["key"]
             if state_key.shape == (624,) and state_key.dtype == np.uint32:
-                return True
+                # TODO: Add an option to convert to a `RandomState` instance?
+                return data
 
-        return False
+        raise TypeError()
 
     @staticmethod
     def values_eq(a, b):
@@ -114,7 +122,7 @@ aesara.compile.register_view_op_c_code(
 random_state_type = RandomStateType()
 
 
-class RandomGeneratorType(RandomType):
+class RandomGeneratorType(RandomType[np.random.Generator]):
     r"""A Type wrapper for `numpy.random.Generator`.
 
     The reason this exists (and `Generic` doesn't suffice) is that
@@ -130,16 +138,25 @@ class RandomGeneratorType(RandomType):
     def __repr__(self):
         return "RandomGeneratorType"
 
-    @staticmethod
-    def is_valid_value(a, strict):
-        if isinstance(a, np.random.Generator):
-            return True
+    def filter(self, data, strict=False, allow_downcast=None):
+        """
+        XXX: This doesn't convert `data` to the same type of underlying RNG type
+        as `self`.  It really only checks that `data` is of the appropriate type
+        to be a valid `RandomGeneratorType`.
+
+        In other words, it serves as a `Type.is_valid_value` implementation,
+        but, because the default `Type.is_valid_value` depends on
+        `Type.filter`, we need to have it here to avoid surprising circular
+        dependencies in sub-classes.
+        """
+        if isinstance(data, np.random.Generator):
+            return data
 
-        if not strict and isinstance(a, dict):
-            if "bit_generator" not in a:
-                return False
+        if not strict and isinstance(data, dict):
+            if "bit_generator" not in data:
+                raise TypeError()
             else:
-                bit_gen_key = a["bit_generator"]
+                bit_gen_key = data["bit_generator"]
 
                 if hasattr(bit_gen_key, "_value"):
                     bit_gen_key = int(bit_gen_key._value)
@@ -148,16 +165,16 @@ class RandomGeneratorType(RandomType):
                 gen_keys, state_keys = gen_states_keys[bit_gen_key]
 
                 for key in gen_keys:
-                    if key not in a:
-                        return False
+                    if key not in data:
+                        raise TypeError()
 
                 for key in state_keys:
-                    if key not in a["state"]:
-                        return False
+                    if key not in data["state"]:
+                        raise TypeError()
 
-                return True
+                return data
 
-        return False
+        raise TypeError()
 
     @staticmethod
     def values_eq(a, b):


Test patch:
diff --git a/tests/tensor/random/test_type.py b/tests/tensor/random/test_type.py
index 5f45af2fe..a34d4cbad 100644
--- a/tests/tensor/random/test_type.py
+++ b/tests/tensor/random/test_type.py
@@ -56,15 +56,17 @@ class TestRandomStateType:
         with pytest.raises(TypeError):
             rng_type.filter(1)
 
-        rng = rng.get_state(legacy=False)
-        assert rng_type.is_valid_value(rng, strict=False)
+        rng_dict = rng.get_state(legacy=False)
 
-        rng["state"] = {}
+        assert rng_type.is_valid_value(rng_dict) is False
+        assert rng_type.is_valid_value(rng_dict, strict=False)
 
-        assert rng_type.is_valid_value(rng, strict=False) is False
+        rng_dict["state"] = {}
 
-        rng = {}
-        assert rng_type.is_valid_value(rng, strict=False) is False
+        assert rng_type.is_valid_value(rng_dict, strict=False) is False
+
+        rng_dict = {}
+        assert rng_type.is_valid_value(rng_dict, strict=False) is False
 
     def test_values_eq(self):
 
@@ -147,15 +149,17 @@ class TestRandomGeneratorType:
         with pytest.raises(TypeError):
             rng_type.filter(1)
 
-        rng = rng.__getstate__()
-        assert rng_type.is_valid_value(rng, strict=False)
+        rng_dict = rng.__getstate__()
+
+        assert rng_type.is_valid_value(rng_dict) is False
+        assert rng_type.is_valid_value(rng_dict, strict=False)
 
-        rng["state"] = {}
+        rng_dict["state"] = {}
 
-        assert rng_type.is_valid_value(rng, strict=False) is False
+        assert rng_type.is_valid_value(rng_dict, strict=False) is False
 
-        rng = {}
-        assert rng_type.is_valid_value(rng, strict=False) is False
+        rng_dict = {}
+        assert rng_type.is_valid_value(rng_dict, strict=False) is False
 
     def test_values_eq(self):