make inner and _stein_metric support keepdim

geoopt/linalg/batch_linalg.py

@@ -1,3 +1,5 @@
+from typing import List, Callable, Tuple
+import torch
 import torch.jit
 from . import _expm
 
@@ -91,6 +93,23 @@ def block_matrix(blocks: List[List[torch.Tensor]], dim0: int = -2, dim1: int = -
     return torch.cat(hblocks, dim=dim0)
 
 
+@torch.jit.script
+def trace(x: torch.Tensor) -> torch.Tensor:
+    r"""self-implemented matrix trace, since `torch.trace` only support 2-d input.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        input matrix
+
+    Returns
+    -------
+    torch.Tensor
+        :math:`\operationname{Tr}(x)`
+    """
+    return torch.diagonal(x, dim1=-2, dim2=-1).sum(-1)
+
+
 def sym_funcm(
     x: torch.Tensor, func: Callable[[torch.Tensor], torch.Tensor]
 ) -> torch.Tensor:

geoopt/manifolds/symmetric_positive_definite.py

@@ -99,9 +99,10 @@ def _stein_metric(
         def log_det(tensor: torch.Tensor) -> torch.Tensor:
             return torch.log(torch.det(tensor))
 
+        ret = log_det((x + y) * 0.5) - 0.5 * log_det(x @ y)
         if keepdim:
-            raise ValueError("`torch.det` doesn't support keepdim.")
-        return log_det((x + y) * 0.5) - 0.5 * log_det(x @ y)
+            return torch.unsqueeze(torch.unsqueeze(ret, -1), -1)
+        return ret
 
     def _log_eucliden_metric(
         self, x: torch.Tensor, y: torch.Tensor, keepdim=False
@@ -231,10 +232,11 @@ def inner(
         """
         if v is None:
             v = u
-        if keepdim:
-            raise ValueError("`torch.trace` doesn't support keepdim.")
         inv_x = batch_linalg.sym_invm(x)
-        return torch.trace(inv_x @ u @ inv_x @ v)
+        ret = batch_linalg.trace(inv_x @ u @ inv_x @ v)
+        if keepdim:
+            return torch.unsqueeze(torch.unsqueeze(ret, -1), -1)
+        return ret
 
     def retr(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
         inv_x = batch_linalg.sym_invm(x)