Avoid recompiling [d]logp in every step of find_MAP

This also fixes a bug where graphs with NotImplemented gradients were not detected before deciding on optimization method

Author: ricardoV94

pymc/tests/test_starting.py

@@ -46,7 +46,7 @@ def test_accuracy_non_normal():
         close_to(newstart["x"], mu, select_by_precision(float64=1e-5, float32=1e-4))
 
 
-@pytest.mark.xfail(reason="find_MAP fails with derivatives")
+@pytest.mark.xfail(reason="first call to find_MAP is failing")
 def test_find_MAP_discrete():
     tol = 2.0 ** -11
     alpha = 4
@@ -68,15 +68,12 @@ def test_find_MAP_discrete():
     assert map_est2["ss"] == 14
 
 
-@pytest.mark.xfail(reason="find_MAP fails with derivatives")
 def test_find_MAP_no_gradient():
     _, model = simple_arbitrary_det()
     with model:
         find_MAP()
 
 
-@pytest.mark.skip(reason="test is slow because it's failing")
-@pytest.mark.xfail(reason="find_MAP fails with derivatives")
 def test_find_MAP():
     tol = 2.0 ** -11  # 16 bit machine epsilon, a low bar
     data = np.random.randn(100)

pymc/tuning/starting.py

@@ -113,21 +113,15 @@ def find_MAP(
 
     # TODO: If the mapping is fixed, we can simply create graphs for the
     # mapping and avoid all this bijection overhead
-    def logp_func(x):
-        return DictToArrayBijection.mapf(model.compile_logp(jacobian=False))(
-            RaveledVars(x, x0.point_map_info)
-        )
+    compiled_logp_func = DictToArrayBijection.mapf(model.compile_logp(jacobian=False))
+    logp_func = lambda x: compiled_logp_func(RaveledVars(x, x0.point_map_info))
 
     rvs = [model.values_to_rvs[value] for value in vars]
     try:
         # This might be needed for calls to `dlogp_func`
         # start_map_info = tuple((v.name, v.shape, v.dtype) for v in vars)
-
-        def dlogp_func(x):
-            return DictToArrayBijection.mapf(model.compile_dlogp(rvs, jacobian=False))(
-                RaveledVars(x, x0.point_map_info)
-            )
-
+        compiled_dlogp_func = DictToArrayBijection.mapf(model.compile_dlogp(rvs, jacobian=False))
+        dlogp_func = lambda x: compiled_dlogp_func(RaveledVars(x, x0.point_map_info))
         compute_gradient = True
     except (AttributeError, NotImplementedError, tg.NullTypeGradError):
         compute_gradient = False