Revert "Remove BLAS.asum(x) rules"

This reverts commit 3fe59b4.
JuliaDiff · Jan 12, 2020 · 63b9f7e · 63b9f7e
1 parent 61b425b
commit 63b9f7e
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/src/rulesets/LinearAlgebra/blas.jl b/src/rulesets/LinearAlgebra/blas.jl
@@ -67,6 +67,20 @@ end
 ##### `BLAS.asum`
 #####
 
+function frule(::typeof(BLAS.asum), x, _, Δx)
+    return BLAS.asum(x), sum(zip(x, Δx)) do xs
+        x, Δx = xs
+        return sign(x) * Δx
+    end
+end
+
+function rrule(::typeof(BLAS.asum), x)
+    function asum_pullback(ΔΩ)
+        return (NO_FIELDS, @thunk(ΔΩ * sign.(x)))
+    end
+    return BLAS.asum(x), asum_pullback
+end
+
 function rrule(::typeof(BLAS.asum), n, X, incx)
     Ω = BLAS.asum(n, X, incx)
     function asum_pullback(ΔΩ)