Zygote workarounds

baggepinnen · Jun 26, 2020 · 8fc5e74 · 8fc5e74
1 parent 1b708f7
commit 8fc5e74
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/src/utils.jl b/src/utils.jl
@@ -30,12 +30,20 @@ end
 normalize x sums to 1
 """
 function s1(x, dims=:)
-    m = minimum(x, dims=dims)
-    x = copy(x)
-    if any(<(0), m)
-        x .-= m
+    if isderiving()
+        m = minimum(real(x), dims=dims)
+        if any(<(0), m)
+            x = x .- m
+        end
+        return x ./ sum(x, dims=dims)
+    else
+        m = minimum(x, dims=dims)
+        x = float.(x)
+        if any(<(0), m)
+            x .-= m
+        end
+        x ./= sum(x, dims=dims)
     end
-    x ./= sum(x, dims=dims)
 end
 
 """

diff --git a/test/test_barycenter.jl b/test/test_barycenter.jl
@@ -209,7 +209,7 @@ end
 
 
     a = ones(k) |> s1
-    b = [[1,2,3,4] |> s1 for _ in eachindex(Y)]
+    b = [[1.0,2,3,4] |> s1 for _ in eachindex(Y)]
     M = SpectralDistances.distmat_euclidean(X,Y[1])
     g1,a1,b1 = ot_jump(M,a,b[1]) .|> r6
     g2,a2,b2 = sinkhorn_log(M,a,b[1], β=0.001, iters=50000, printerval=5000, tol=1e-9) .|> r6