Group compile_pymc tests in own class

Author: ricardoV94

pymc/tests/test_aesaraf.py

@@ -596,91 +596,88 @@ def test_rvs_to_value_vars_nested():
         assert equal_computations(before, after)
 
 
-def test_check_bounds_flag():
-    """Test that CheckParameterValue Ops are replaced or removed when using compile_pymc"""
-    logp = at.ones(3)
-    cond = np.array([1, 0, 1])
-    bound = check_parameters(logp, cond)
-
-    with pm.Model() as m:
-        pass
-
-    with pytest.raises(ParameterValueError):
-        aesara.function([], bound)()
-
-    m.check_bounds = False
-    with m:
-        assert np.all(compile_pymc([], bound)() == 1)
-
-    m.check_bounds = True
-    with m:
-        assert np.all(compile_pymc([], bound)() == -np.inf)
-
-
-def test_compile_pymc_sets_rng_updates():
-    rng = aesara.shared(np.random.default_rng(0))
-    x = pm.Normal.dist(rng=rng)
-    assert x.owner.inputs[0] is rng
-    f = compile_pymc([], x)
-    assert not np.isclose(f(), f())
-
-    # Check that update was not done inplace
-    assert not hasattr(rng, "default_update")
-    f = aesara.function([], x)
-    assert f() == f()
-
-
-def test_compile_pymc_with_updates():
-    x = aesara.shared(0)
-    f = compile_pymc([], x, updates={x: x + 1})
-    assert f() == 0
-    assert f() == 1
-
-
-def test_compile_pymc_missing_default_explicit_updates():
-    rng = aesara.shared(np.random.default_rng(0))
-    x = pm.Normal.dist(rng=rng)
-
-    # By default, compile_pymc should update the rng of x
-    f = compile_pymc([], x)
-    assert f() != f()
-
-    # An explicit update should override the default_update, like aesara.function does
-    # For testing purposes, we use an update that leaves the rng unchanged
-    f = compile_pymc([], x, updates={rng: rng})
-    assert f() == f()
-
-    # If we specify a custom default_update directly it should use that instead.
-    rng.default_update = rng
-    f = compile_pymc([], x)
-    assert f() == f()
-
-    # And again, it should be overridden by an explicit update
-    f = compile_pymc([], x, updates={rng: x.owner.outputs[0]})
-    assert f() != f()
-
-
-def test_compile_pymc_updates_inputs():
-    """Test that compile_pymc does not include rngs updates of variables that are inputs
-    or ancestors to inputs
-    """
-    x = at.random.normal()
-    y = at.random.normal(x)
-    z = at.random.normal(y)
-
-    for inputs, rvs_in_graph in (
-        ([], 3),
-        ([x], 2),
-        ([y], 1),
-        ([z], 0),
-        ([x, y], 1),
-        ([x, y, z], 0),
-    ):
-        fn = compile_pymc(inputs, z, on_unused_input="ignore")
-        fn_fgraph = fn.maker.fgraph
-        # Each RV adds a shared input for its rng
-        assert len(fn_fgraph.inputs) == len(inputs) + rvs_in_graph
-        # If the output is an input, the graph has a DeepCopyOp
-        assert len(fn_fgraph.apply_nodes) == max(rvs_in_graph, 1)
-        # Each RV adds a shared output for its rng
-        assert len(fn_fgraph.outputs) == 1 + rvs_in_graph
+class TestCompilePyMC:
+    def test_check_bounds_flag(self):
+        """Test that CheckParameterValue Ops are replaced or removed when using compile_pymc"""
+        logp = at.ones(3)
+        cond = np.array([1, 0, 1])
+        bound = check_parameters(logp, cond)
+
+        with pm.Model() as m:
+            pass
+
+        with pytest.raises(ParameterValueError):
+            aesara.function([], bound)()
+
+        m.check_bounds = False
+        with m:
+            assert np.all(compile_pymc([], bound)() == 1)
+
+        m.check_bounds = True
+        with m:
+            assert np.all(compile_pymc([], bound)() == -np.inf)
+
+    def test_compile_pymc_sets_rng_updates(self):
+        rng = aesara.shared(np.random.default_rng(0))
+        x = pm.Normal.dist(rng=rng)
+        assert x.owner.inputs[0] is rng
+        f = compile_pymc([], x)
+        assert not np.isclose(f(), f())
+
+        # Check that update was not done inplace
+        assert not hasattr(rng, "default_update")
+        f = aesara.function([], x)
+        assert f() == f()
+
+    def test_compile_pymc_with_updates(self):
+        x = aesara.shared(0)
+        f = compile_pymc([], x, updates={x: x + 1})
+        assert f() == 0
+        assert f() == 1
+
+    def test_compile_pymc_missing_default_explicit_updates(self):
+        rng = aesara.shared(np.random.default_rng(0))
+        x = pm.Normal.dist(rng=rng)
+
+        # By default, compile_pymc should update the rng of x
+        f = compile_pymc([], x)
+        assert f() != f()
+
+        # An explicit update should override the default_update, like aesara.function does
+        # For testing purposes, we use an update that leaves the rng unchanged
+        f = compile_pymc([], x, updates={rng: rng})
+        assert f() == f()
+
+        # If we specify a custom default_update directly it should use that instead.
+        rng.default_update = rng
+        f = compile_pymc([], x)
+        assert f() == f()
+
+        # And again, it should be overridden by an explicit update
+        f = compile_pymc([], x, updates={rng: x.owner.outputs[0]})
+        assert f() != f()
+
+    def test_compile_pymc_updates_inputs(self):
+        """Test that compile_pymc does not include rngs updates of variables that are inputs
+        or ancestors to inputs
+        """
+        x = at.random.normal()
+        y = at.random.normal(x)
+        z = at.random.normal(y)
+
+        for inputs, rvs_in_graph in (
+            ([], 3),
+            ([x], 2),
+            ([y], 1),
+            ([z], 0),
+            ([x, y], 1),
+            ([x, y, z], 0),
+        ):
+            fn = compile_pymc(inputs, z, on_unused_input="ignore")
+            fn_fgraph = fn.maker.fgraph
+            # Each RV adds a shared input for its rng
+            assert len(fn_fgraph.inputs) == len(inputs) + rvs_in_graph
+            # If the output is an input, the graph has a DeepCopyOp
+            assert len(fn_fgraph.apply_nodes) == max(rvs_in_graph, 1)
+            # Each RV adds a shared output for its rng
+            assert len(fn_fgraph.outputs) == 1 + rvs_in_graph