scalig

geoopt · Nov 30, 2019 · fdeab42 · fdeab42
1 parent f1b0e69
commit fdeab42
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/geoopt/manifolds/base.py b/geoopt/manifolds/base.py
@@ -81,9 +81,11 @@ def mobius_fn_apply(
                     return res
     """
 
-    def __call__(self, scaling_info: ScalingInfo):
+    def __call__(self, scaling_info: ScalingInfo, *aliases):
         def register(fn):
             self[fn.__name__] = scaling_info
+            for alias in aliases:
+                self[alias] = scaling_info
             return fn
 
         return register

diff --git a/geoopt/manifolds/euclidean.py b/geoopt/manifolds/euclidean.py
@@ -107,7 +107,7 @@ def transp(self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor) -> torch.Ten
         target_shape = broadcast_shapes(x.shape, y.shape, v.shape)
         return v.expand(target_shape)
 
-    @__scaling__(ScalingInfo(std=-1))
+    @__scaling__(ScalingInfo(std=-1), "random")
     def random_normal(
         self, *size, mean=0.0, std=1.0, device=None, dtype=None
     ) -> "geoopt.ManifoldTensor":

diff --git a/geoopt/manifolds/poincare/__init__.py b/geoopt/manifolds/poincare/__init__.py
@@ -286,7 +286,7 @@ def mobius_fn_apply_chain(
         else:
             return res
 
-    @__scaling__(ScalingInfo(std=-1))
+    @__scaling__(ScalingInfo(std=-1), "random")
     def random_normal(
         self, *size, mean=0, std=1, dtype=None, device=None
     ) -> "geoopt.ManifoldTensor":