logprob for maximum derived

Author: Dhruvanshu-Joshi

tests/logprob/test_checks.py

@@ -33,6 +33,7 @@
 #   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
+import re
 
 import numpy as np
 import pytensor
@@ -43,7 +44,7 @@
 from scipy import stats
 
 from pymc.distributions import Dirichlet
-from pymc.logprob.joint_logprob import factorized_joint_logprob
+from pymc.logprob.basic import conditional_logp
 from tests.distributions.test_multivariate import dirichlet_logpdf
 
 
@@ -58,7 +59,7 @@ def test_specify_shape_logprob():
 
     # 2. Request logp
     x_vv = x_rv.clone()
-    [x_logp] = factorized_joint_logprob({x_rv: x_vv}).values()
+    [x_logp] = conditional_logp({x_rv: x_vv}).values()
 
     # 3. Test logp
     x_logp_fn = pytensor.function([last_dim, x_vv], x_logp)
@@ -80,17 +81,19 @@ def test_assert_logprob():
     rv = pt.random.normal()
     assert_op = Assert("Test assert")
     # Example: Add assert that rv must be positive
-    assert_rv = assert_op(rv > 0, rv)
+    assert_rv = assert_op(rv, rv > 0)
     assert_rv.name = "assert_rv"
 
     assert_vv = assert_rv.clone()
-    assert_logp = factorized_joint_logprob({assert_rv: assert_vv})[assert_vv]
+    assert_logp = conditional_logp({assert_rv: assert_vv})[assert_vv]
 
     # Check valid value is correct and doesn't raise
     # Since here the value to the rv satisfies the condition, no error is raised.
     valid_value = 3.0
-    with pytest.raises(AssertionError, match="Test assert"):
-        assert_logp.eval({assert_vv: valid_value})
+    np.testing.assert_allclose(
+        assert_logp.eval({assert_vv: valid_value}),
+        stats.norm.logpdf(valid_value),
+    )
 
     # Check invalid value
     # Since here the value to the rv is negative, an exception is raised as the condition is not met