Fix "auto" for forward-mode vs reverse-mode jacobians

brentyi · Dec 23, 2024 · 0991ee8 · 0991ee8
1 parent 67239c1
commit 0991ee8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/jaxls/_factor_graph.py b/src/jaxls/_factor_graph.py
@@ -134,7 +134,7 @@ def compute_jac_with_perturb(factor: _AnalyzedFactor) -> jax.Array:
                     "reverse": jax.jacrev,
                     "auto": jax.jacrev
                     if factor.residual_dim < val_subset._get_tangent_dim()
-                    else jax.jacrev,
+                    else jax.jacfwd,
                 }[factor.jac_mode]
                 return jacfunc(
                     # The residual function, with respect to to some local delta.