Combine with ForwardDiff: `ERROR: InexactError: trunc(Int64, NaN)`

[ranocha]

Currently, Octavian.jl does not seem to support AD using ForwardDiff.jl:

julia> using Pkg; Pkg.activate(temp=true); Pkg.add(["ForwardDiff", "Octavian"])
[...]
  [f6369f11] + ForwardDiff v0.10.21
  [6fd5a793] + Octavian v0.3.4
[...]

julia> using ForwardDiff, Octavian, LinearAlgebra

julia> u = rand(50, 50); du = similar(u); A = rand(size(u, 1), size(u, 1));

julia> config = ForwardDiff.JacobianConfig(nothing, du, u);

julia> ForwardDiff.jacobian((du, u) -> mul!(du, A, u), du, u, config) ≈ kron(I(size(u, 2)), A) # stdlib is fine
true

julia> ForwardDiff.jacobian((du, u) -> matmul!(du, A, u), du, u, config) ≈ kron(I(size(u, 2)), A)
ERROR: InexactError: trunc(Int64, NaN)
Stacktrace:
  [1] trunc
    @ ./float.jl:805 [inlined]
  [2] floor(#unused#::Type{Int64}, x::Float64)
    @ Base ./float.jl:367
  [3] #s3#19
    @ ~/.julia/packages/Static/A5kpr/src/float.jl:106 [inlined]
  [4] var"#s3#19"(M::Any, ::Any, #unused#::Any)
    @ Static ./none:0
  [5] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:580
  [6] block_sizes(#unused#::Val{ForwardDiff.Dual{Nothing, Float64, 12}}, W::StaticInt{0}, α::Static.StaticFloat64{0.0}, β::Static.StaticFloat64{0.0}, L₁ₑ::Static.StaticFloat64{1353.3202323303008}, L₂ₑ::Static.StaticFloat64{66157.2550994403})
    @ Octavian ~/.julia/packages/Octavian/LB784/src/block_sizes.jl:16
  [7] block_sizes
    @ ~/.julia/packages/Octavian/LB784/src/block_sizes.jl:10 [inlined]
  [8] block_sizes
    @ ~/.julia/packages/Octavian/LB784/src/block_sizes.jl:23 [inlined]
  [9] __matmul!(C::LayoutPointers.StridedPointer{ForwardDiff.Dual{Nothing, Float64, 12}, 2, 1, 0, (1, 2), Tuple{StaticInt{104}, Int64}, Tuple{StaticInt{0}, StaticInt{0}}}, A::LayoutPointers.StridedPointer{Float64, 2, 1, 0, (1, 2), Tuple{StaticInt{8}, Int64}, Tuple{StaticInt{0}, StaticInt{0}}}, B::LayoutPointers.StridedPointer{ForwardDiff.Dual{Nothing, Float64, 12}, 2, 1, 0, (1, 2), Tuple{StaticInt{104}, Int64}, Tuple{StaticInt{0}, StaticInt{0}}}, α::StaticInt{1}, β::StaticInt{0}, M::Int64, K::Int64, N::Int64, nthread::Nothing)
    @ Octavian ~/.julia/packages/Octavian/LB784/src/matmul.jl:343
 [10] _matmul!
    @ ~/.julia/packages/Octavian/LB784/src/matmul.jl:289 [inlined]
 [11] matmul!
    @ ~/.julia/packages/Octavian/LB784/src/matmul.jl:256 [inlined]
 [12] matmul!(C::Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}, A::Matrix{Float64}, B::Matrix{ForwardDiff.Dual{Nothing, Float64, 12}})
    @ Octavian ~/.julia/packages/Octavian/LB784/src/matmul.jl:236
 [13] (::var"#9#10")(du::Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}, u::Matrix{ForwardDiff.Dual{Nothing, Float64, 12}})
    @ Main ./REPL[9]:1
 [14] chunk_mode_jacobian(f!::var"#9#10", y::Matrix{Float64}, x::Matrix{Float64}, cfg::ForwardDiff.JacobianConfig{Nothing, Float64, 12, Tuple{Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}, Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UDrkY/src/jacobian.jl:223
 [15] jacobian(f!::Function, y::Matrix{Float64}, x::Matrix{Float64}, cfg::ForwardDiff.JacobianConfig{Nothing, Float64, 12, Tuple{Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}, Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}}}, ::Val{true})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UDrkY/src/jacobian.jl:40
 [16] jacobian(f!::Function, y::Matrix{Float64}, x::Matrix{Float64}, cfg::ForwardDiff.JacobianConfig{Nothing, Float64, 12, Tuple{Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}, Matrix{ForwardDiff.Dual{Nothing, Float64, 12}}}})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UDrkY/src/jacobian.jl:36
 [17] top-level scope
    @ REPL[9]:1

I'm using Julia version 1.7.0-rc1.

[chriselrod]

We could add dual-number kernels.
If only one input contains dual, it'd be a pretty trival reinterpret.
If both inputs have duals, you could do two successive matmuls, the second using a view to cutoff the extra part.