Use model context in logp_dlogp_function to respect check_bounds

Author: ricardoV94

pymc/model/core.py

@@ -564,15 +564,16 @@ def logp_dlogp_function(
             for var in self.value_vars
             if var in input_vars and var not in grad_vars
         }
-        return ValueGradFunction(
-            costs,
-            grad_vars,
-            extra_vars_and_values,
-            model=self,
-            initial_point=initial_point,
-            ravel_inputs=ravel_inputs,
-            **kwargs,
-        )
+        with self:
+            return ValueGradFunction(
+                costs,
+                grad_vars,
+                extra_vars_and_values,
+                model=self,
+                initial_point=initial_point,
+                ravel_inputs=ravel_inputs,
+                **kwargs,
+            )
 
     def compile_logp(
         self,

tests/model/test_core.py

@@ -443,6 +443,15 @@ def test_missing_data(self):
         # Assert that all the elements of res are equal
         assert res[1:] == res[:-1]
 
+    def test_check_bounds_out_of_model_context(self):
+        with pm.Model(check_bounds=False) as m:
+            x = pm.Normal("x")
+            y = pm.Normal("y", sigma=x)
+        fn = m.logp_dlogp_function(ravel_inputs=True)
+        fn.set_extra_values({})
+        # When there are no bounds check logp turns into `nan`
+        assert np.isnan(fn(np.array([-1.0, -1.0]))[0])
+
 
 class TestPytensorRelatedLogpBugs:
     def test_pytensor_switch_broadcast_edge_cases_1(self):