pymc-extras 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

tests/statespace/test_structural.py

@@ -2,6 +2,7 @@ import functools as ft
 import warnings
 
 from collections import defaultdict
+from copyreg import remove_extension
 from typing import Optional
 
 import numpy as np
@@ -592,13 +593,18 @@ def test_autoregressive_model(order, rng):
 
 @pytest.mark.parametrize("s", [10, 25, 50])
 @pytest.mark.parametrize("innovations", [True, False])
-def test_time_seasonality(s, innovations, rng):
+@pytest.mark.parametrize("remove_first_state", [True, False])
+def test_time_seasonality(s, innovations, remove_first_state, rng):
     def random_word(rng):
         return "".join(rng.choice(list("abcdefghijklmnopqrstuvwxyz")) for _ in range(5))
 
     state_names = [random_word(rng) for _ in range(s)]
     mod = st.TimeSeasonality(
-        season_length=s, innovations=innovations, name="season", state_names=state_names
+        season_length=s,
+        innovations=innovations,
+        name="season",
+        state_names=state_names,
+        remove_first_state=remove_first_state,
     )
     x0 = np.zeros(mod.k_states, dtype=floatX)
     x0[0] = 1
@@ -615,7 +621,8 @@ def test_time_seasonality(s, innovations, rng):
     # Check coords
     mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
-    assert mod.coords["season_state"] == state_names[1:]
+    test_slice = slice(1, None) if remove_first_state else slice(None)
+    assert mod.coords["season_state"] == state_names[test_slice]
 
 
 def get_shift_factor(s):

tests/test_pivoted_cholesky.py

@@ -8,7 +8,7 @@
 #     pass
 # import numpy as np
 #
-# import pymc_experimental as pmx
+# import pymc_extras as pmx
 #
 #
 # def test_match_gpytorch_linearcg_output():

tests/utils.py

@@ -1,31 +0,0 @@
-from collections.abc import Sequence
-
-from pytensor.compile import SharedVariable
-from pytensor.graph import Constant, graph_inputs
-from pytensor.graph.basic import Variable, equal_computations
-from pytensor.tensor.random.type import RandomType
-
-
-def equal_computations_up_to_root(
-    xs: Sequence[Variable], ys: Sequence[Variable], ignore_rng_values=True
-) -> bool:
-    # Check if graphs are equivalent even if root variables have distinct identities
-
-    x_graph_inputs = [var for var in graph_inputs(xs) if not isinstance(var, Constant)]
-    y_graph_inputs = [var for var in graph_inputs(ys) if not isinstance(var, Constant)]
-    if len(x_graph_inputs) != len(y_graph_inputs):
-        return False
-    for x, y in zip(x_graph_inputs, y_graph_inputs):
-        if x.type != y.type:
-            return False
-        if x.name != y.name:
-            return False
-        if isinstance(x, SharedVariable):
-            if not isinstance(y, SharedVariable):
-                return False
-            if isinstance(x.type, RandomType) and ignore_rng_values:
-                continue
-            if not x.type.values_eq(x.get_value(), y.get_value()):
-                return False
-
-    return equal_computations(xs, ys, in_xs=x_graph_inputs, in_ys=y_graph_inputs)