bugfix

Author: nocotan

VERSION

@@ -1 +1 @@
-1.0.6
+1.0.7

gs_divergence/__init__.py

@@ -1,4 +1,4 @@
 from gs_divergence.geodesical_skew_divergence import gs_div
 from gs_divergence.symmetrized_geodesical_skew_divergence import symmetrized_gs_div
 
-__version__ = '1.0.6'
+__version__ = '1.0.7'

gs_divergence/alpha_geodesic.py

@@ -16,16 +16,16 @@ def alpha_geodesic(
     if alpha == 1:
         return torch.exp((1 - lmd) * torch.log(a_) + lmd * torch.log(b_))
     elif alpha >= 1e+9:
-        return torch.min(a, b)
+        return torch.min(a_, b_)
     elif alpha <= -1e+9:
-        return torch.max(a, b)
+        return torch.max(a_, b_)
     else:
         p = (1 - alpha) / 2
         lhs = a_ ** p
         rhs = b_ ** p
         g = ((1 - lmd) * lhs + lmd * rhs) ** (1/p)
 
         if alpha > 0 and (g == 0).sum() > 0:
-            return torch.min(a, b)
+            return torch.min(a_, b_)
 
         return g

tests/test_alpha_geodesic.py

@@ -74,7 +74,7 @@ def test_value_inf(self):
         g = alpha_geodesic(a, b, alpha=100, lmd=0.5)
         res = torch.min(a, b)
 
-        self.assertTrue(torch.equal(g, res))
+        self.assertTrue(torch.all(torch.isclose(g, res)))
 
     def test_grad(self):
         a = torch.tensor([[0.1, 0.2, 0.7], [0.5, 0.5, 0.0]], requires_grad=True)