JuliaDecisionFocusedLearning / ImplicitDifferentiation.jl

Automatic differentiation of implicit functions
https://juliadecisionfocusedlearning.github.io/ImplicitDifferentiation.jl/
MIT License
123 stars 6 forks source link

Use `block_gmres` for batched linear solves #129

Open thorek1 opened 11 months ago

thorek1 commented 11 months ago

The latest Krylov.jl version has a block-gmres version which I think can simplify/speed-up some parts of the code.

gdalle commented 10 months ago

Thanks for the idea. I don't have much time to work on this at the moment though

mohamed82008 commented 10 months ago

how can block-gmres help ID.jl?

thorek1 commented 10 months ago

for example here:

    dy = ntuple(Val(N)) do k
        dₖx = partials.(x_and_dx, k)
        dₖx_vec = vec(dₖx)
        dₖc_vec = similar(y_vec)
        mul!(dₖc_vec, B_vec, dₖx_vec)
        dₖy_vec = solve(implicit.linear_solver, A_vec_presolved, -dₖc_vec)
        reshape(dₖy_vec, size(y))
    end

    y_and_dy = map(eachindex(IndexCartesian(), y)) do i
        Dual{T}(y[i], Partials(ntuple(k -> dy[k][i], Val(N))))
    end

instead of solving with gmres for each slice you could use block-gmres once. I didn't check the block-gmres implementation but I think it is worth testing whether this speeds up computation and this might also allow to solve #118

gdalle commented 6 months ago

This is implemented in #135 but it still requires https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/854 to be fixed