Refactor batched_vec (#464)

Branching on the type of the second argument caused a subtle performance bug when differentiating via `Zygote`; see #462

src/batched/batchedmul.jl

@@ -172,13 +172,12 @@ julia> batched_vec(A,b) |> size
 (16, 32)
 ```
 """
-function batched_vec(A::AbstractArray{T,3} where T, B::AbstractMatrix)
- # If B is transposed, then stride=1 is the batch dim, so we will end up copying anyway:
- if B isa AdjOrTransAbsMat{<:BlasFloat, <:StridedMatrix}
- return batched_vec(A, copy(B))
- end
+batched_vec(A::AbstractArray{T,3} where T, B::AbstractMatrix) =
  reshape(batched_mul(A, reshape(B, size(B,1), 1, size(B,2))), size(A,1), size(A,3))
-end
+
+# If B is transposed, then stride=1 is the batch dim, so we will end up copying anyway:
+batched_vec(A::AbstractArray{T,3} where T, B::AdjOrTransAbsMat{<:BlasFloat, <:StridedMatrix}) =
+ batched_vec(A, copy(B))
 
 batched_vec(A::AbstractArray{T,3} where T, b::AbstractVector) =
  reshape(batched_mul(A, reshape(b, length(b), 1, 1)), size(A,1), size(A,3))