Merge 44e3fce into 03a2542

geoopt · Nov 14, 2019 · 283b568 · 283b568
2 parents 03a2542 + 44e3fce
commit 283b568
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)