fix fix for #22

SciML · Apr 2, 2020 · b969b9f · b969b9f
1 parent eb9d499
commit b969b9f
Showing 1 changed file with 10 additions and 13 deletions.
diff --git a/src/blas.jl b/src/blas.jl
@@ -5,15 +5,17 @@ _name(x::Expr) = (@assert x.head == :(::); x.args[1])
 macro reroute(f, g)
     fname = f.args[1]
     fargs = f.args[2:end]
+    gname = g.args[1]
+    gargs = g.args[2:end]
     quote
         @inline function Cassette.overdub(ctx::SparsityContext,
                                           f::typeof($(esc(fname))),
                                           $(fargs...))
             Cassette.recurse(
                 ctx,
                 invoke,
-                f,
-                $(esc(:(Tuple{$(g.args[2:end]...)}))),
+                $(esc(gname)),
+                $(esc(:(Tuple{$(gargs...)}))),
                 $(map(_name, fargs)...))
         end
 
@@ -23,8 +25,8 @@ macro reroute(f, g)
             Cassette.recurse(
                 ctx,
                 invoke,
-                $(esc(g.args[1])),
-                $(esc(:(Tuple{$(g.args[2:end]...)}))),
+                $(esc(gname)),
+                $(esc(:(Tuple{$(gargs...)}))),
                 $(map(_name, fargs)...))
         end
     end
@@ -34,12 +36,7 @@ end
 @reroute LinearAlgebra.BLAS.axpy!(x, y) LinearAlgebra.axpy!(Any,
                                                       AbstractArray,
                                                       AbstractArray)
-@reroute LinearAlgebra.mul!(y::AbstractVector,
-                            A::AbstractVecOrMat,
-                            x::AbstractVector,
-                            α::Number,
-                            β::Number) LinearAlgebra.mul!(AbstractVector,
-                                                          AbstractVecOrMat,
-                                                          AbstractVector,
-                                                          Number,
-                                                          Number)
+
+gengemv!(tA, α, A, x, β, y) = LinearAlgebra.generic_matvecmul!(y, tA, A, x, LinearAlgebra.MulAddMul(α, β))
+
+@reroute LinearAlgebra.BLAS.gemv!(tA, α, A, x, β, y) gengemv!(Any, Any, Any, Any, Any, Any)