Rename compute_covariance to compute_hessian

Author: jessegrabowski

pymc_extras/inference/laplace_approx/find_map.py

@@ -198,7 +198,7 @@ def find_MAP(
     include_transformed: bool = True,
     gradient_backend: GradientBackend = "pytensor",
     compile_kwargs: dict | None = None,
-    compute_covariance: bool = False,
+    compute_hessian: bool = False,
     **optimizer_kwargs,
 ) -> (
     dict[str, np.ndarray]
@@ -240,7 +240,7 @@ def find_MAP(
         Whether to include transformed variable values in the returned dictionary. Defaults to True.
     gradient_backend: str, default "pytensor"
         Which backend to use to compute gradients. Must be one of "pytensor" or "jax".
-    compute_covariance: bool
+    compute_hessian: bool
         If True, the inverse Hessian matrix at the optimum will be computed and included in the returned
         InferenceData object. This is needed for the Laplace approximation, but can be computationally expensive for
         high-dimensional problems. Defaults to False.
@@ -321,7 +321,7 @@ def find_MAP(
             **optimizer_kwargs,
         )
 
-    if compute_covariance:
+    if compute_hessian:
         H_inv = _compute_inverse_hessian(
             optimizer_result=optimizer_result,
             optimal_point=None,

pymc_extras/inference/laplace_approx/laplace.py

@@ -389,7 +389,7 @@ def fit_laplace(
         include_transformed=include_transformed,
         gradient_backend=gradient_backend,
         compile_kwargs=compile_kwargs,
-        compute_covariance=True,
+        compute_hessian=True,
         **optimizer_kwargs,
     )
 

tests/inference/laplace_approx/test_find_map.py

@@ -133,7 +133,7 @@ def compute_z(x):
     ],
 )
 @pytest.mark.parametrize(
-    "backend, gradient_backend, include_transformed, compute_covariance",
+    "backend, gradient_backend, include_transformed, compute_hessian",
     [("jax", "jax", True, True), ("jax", "pytensor", False, False)],
     ids=str,
 )
@@ -145,7 +145,7 @@ def test_find_MAP(
     backend,
     gradient_backend: GradientBackend,
     include_transformed,
-    compute_covariance,
+    compute_hessian,
     rng,
 ):
     pytest.importorskip("jax")
@@ -165,7 +165,7 @@ def test_find_MAP(
             include_transformed=include_transformed,
             compile_kwargs={"mode": backend.upper()},
             maxiter=5,
-            compute_covariance=compute_covariance,
+            compute_hessian=compute_hessian,
         )
 
     assert hasattr(idata, "posterior")
@@ -186,7 +186,7 @@ def test_find_MAP(
     else:
         assert not hasattr(idata, "unconstrained_posterior")
 
-    assert ("covariance_matrix" in idata.fit) == compute_covariance
+    assert ("covariance_matrix" in idata.fit) == compute_hessian
 
 
 def test_find_map_outside_model_context():