Consistency: project_all() -> normalize_all()

examples/se3_optimization.py

@@ -168,7 +168,7 @@ def step(self: State) -> Tuple[jnp.ndarray, State]:
             new_params = optax.apply_updates(self.params, updates)  # type: ignore
 
             # Project back to manifold.
-            new_params = jaxlie.manifold.project_all(new_params)
+            new_params = jaxlie.manifold.normalize_all(new_params)
 
         elif self.algorithm == "exponential_coordinates":
             # If we parameterize with exponential coordinates, we can

jaxlie/manifold/__init__.py

@@ -1,6 +1,6 @@
 from ._backprop import grad, value_and_grad, zero_tangents
 from ._deltas import rminus, rplus, rplus_jacobian_parameters_wrt_delta
-from ._tree_utils import project_all
+from ._tree_utils import normalize_all
 
 __all__ = [
     "grad",
@@ -9,5 +9,5 @@
     "rminus",
     "rplus",
     "rplus_jacobian_parameters_wrt_delta",
-    "project_all",
+    "normalize_all",
 ]

jaxlie/manifold/_tree_utils.py

@@ -81,7 +81,7 @@ def _map_group_trees(
 PytreeType = TypeVar("PytreeType")
 
 
-def project_all(pytree: PytreeType) -> PytreeType:
+def normalize_all(pytree: PytreeType) -> PytreeType:
     """Call `.normalize()` on each Lie group instance in a pytree.
 
     Results in a naive projection of each group instance to its respective manifold.

tests/test_manifold.py

@@ -87,11 +87,11 @@ def test_rminus_auto_vmap():
     assert_arrays_close(deltas[0], -deltas[1])
 
 
-def test_project():
+def test_normalize():
     container = {"key": (jaxlie.SO3(jnp.array([2.0, 0.0, 0.0, 0.0])),)}
     container_valid = {"key": (jaxlie.SO3(jnp.array([1.0, 0.0, 0.0, 0.0])),)}
     with pytest.raises(AssertionError):
         assert_transforms_close(container["key"][0], container_valid["key"][0])
     assert_transforms_close(
-        jaxlie.manifold.project_all(container)["key"][0], container_valid["key"][0]
+        jaxlie.manifold.normalize_all(container)["key"][0], container_valid["key"][0]
     )