mu.norm does not return ||.||^2!!

ott-jax · Nov 19, 2024 · b2d6973 · b2d6973
1 parent ab0139b
commit b2d6973
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/ott/geometry/costs.py b/src/ott/geometry/costs.py
@@ -363,8 +363,7 @@ def __init__(self, p: float):
     self.p = p
 
   def h(self, z: jnp.ndarray) -> float:  # noqa: D102
-    # Computed by raising squared-norm to p/2.
-    return mu.norm(z) ** (self.p / 2.)
+    return mu.norm(z, ord=2) ** (self.p)
 
   def tree_flatten(self):  # noqa: D102
     return (), (self.p,)