Fix code style

Author: master

tensorflow_riemopt/manifolds/cholesky.py

@@ -47,7 +47,7 @@ def proju(self, x, u):
         u_sym = (utils.transposem(u) + u) / 2.0
         u_diag, u_lower = self._diag_and_strictly_lower(u_sym)
         x_diag = tf.linalg.diag_part(x)
-        return u_lower + tf.linalg.diag(u_diag * x_diag ** 2)
+        return u_lower + tf.linalg.diag(u_diag * x_diag**2)
 
     def inner(self, x, u, v, keepdims=False):
         u_diag, u_lower = self._diag_and_strictly_lower(u)
@@ -92,7 +92,7 @@ def dist(self, x, y, keepdims=False):
             axis=-1,
             keepdims=keepdims,
         )
-        return tf.math.sqrt(lower ** 2 + tf.reshape(diag, lower.shape) ** 2)
+        return tf.math.sqrt(lower**2 + tf.reshape(diag, lower.shape) ** 2)
 
     def ptransp(self, x, y, v):
         x_diag, _ = self._diag_and_strictly_lower(x)

tensorflow_riemopt/manifolds/poincare.py

@@ -55,7 +55,7 @@ def _mobius_add(self, x, y):
         x_y = tf.reduce_sum(x * y, axis=-1, keepdims=True)
         k = tf.cast(self.k, x.dtype)
         return ((1 + 2 * k * x_y + k * y_2) * x + (1 - k * x_2) * y) / (
-            1 + 2 * k * x_y + k ** 2 * x_2 * y_2
+            1 + 2 * k * x_y + k**2 * x_2 * y_2
         )
 
     def _mobius_scal_mul(self, x, r):
@@ -91,15 +91,15 @@ def _lambda(self, x, keepdims=False):
 
     def inner(self, x, u, v, keepdims=False):
         lambda_x = self._lambda(x, keepdims=keepdims)
-        return tf.reduce_sum(u * v, axis=-1, keepdims=keepdims) * lambda_x ** 2
+        return tf.reduce_sum(u * v, axis=-1, keepdims=keepdims) * lambda_x**2
 
     def norm(self, x, u, keepdims=False):
         lambda_x = self._lambda(x, keepdims=keepdims)
         return tf.linalg.norm(u, axis=-1, keepdims=keepdims) * lambda_x
 
     def proju(self, x, u):
         lambda_x = self._lambda(x, keepdims=True)
-        return u / lambda_x ** 2
+        return u / lambda_x**2
 
     def projx(self, x):
         sqrt_k = tf.math.sqrt(tf.cast(self.k, x.dtype))