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

brentyi · Dec 23, 2024 · ee76205 · ee76205
1 parent 67239c1
commit ee76205
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 14 deletions.
diff --git a/.flake8 b/.flake8
diff --git a/.gitignore b/.gitignore
@@ -2,8 +2,3 @@
 *.pyc
 *.egg-info
 __pycache__
-.mypy_cache
-.dmypy.json
-.pytype
-.hypothesis
-.ipynb_checkpoints
diff --git a/LICENSE b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021 Brent Yi
+Copyright (c) 2024 Brent Yi
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

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.