fix poincare dist2

geoopt · Nov 14, 2019 · 44e3fce · 44e3fce
1 parent 03a2542
commit 44e3fce
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/geoopt/manifolds/poincare/__init__.py b/geoopt/manifolds/poincare/__init__.py
@@ -62,6 +62,11 @@ def dist(
     ) -> torch.Tensor:
         return math.dist(x, y, c=self.c, keepdim=keepdim, dim=dim)
 
+    def dist2(
+        self, x: torch.Tensor, y: torch.Tensor, *, keepdim=False, dim=-1
+    ) -> torch.Tensor:
+        return math.dist(x, y, c=self.c, keepdim=keepdim, dim=dim) ** 2
+
     def egrad2rgrad(self, x: torch.Tensor, u: torch.Tensor, *, dim=-1) -> torch.Tensor:
         return math.egrad2rgrad(x, u, c=self.c, dim=dim)