Batched Broyden #43
Batched Broyden #43
Conversation
Codecov Report
@@ Coverage Diff @@
## main #43 +/- ##
==========================================
- Coverage 92.81% 92.39% -0.43%
==========================================
Files 13 14 +1
Lines 557 605 +48
==========================================
+ Hits 517 559 +42
- Misses 40 46 +6
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
Project.toml
Outdated
|
|
||
| [deps] | ||
| ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2" | ||
| DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" | ||
| FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" | ||
| ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" | ||
| LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" | ||
| NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" |
There was a problem hiding this comment.
oh definitely not in this library 😅 Maybe that needs to be an extension.
There was a problem hiding this comment.
😓 I guessed so. Let me set it up as a Pkg Extension then
There was a problem hiding this comment.
For Pkg Extension, I think we will have to set it up as a dispatch on algorithm and not a keyword argument right? Something like Broyden{false} by default and Broyden{true} for the batched version
There was a problem hiding this comment.
yes it probably needs to be split like that. What is the NNlib part for? Why is the matmul batched instead of just a single matmul?
There was a problem hiding this comment.
The matmul is A x B x Batch * B x C x Batch, so they are independent matmuls along the batch dimension. Also IIRC the NNlib implementation does proper CUBLAS dispatches https://github.com/FluxML/NNlibCUDA.jl/blob/5f797aec23cbb5483788697e9e911685b681bbf8/src/batchedmul.jl#L3
|
I made it into an extension but I am not sure why the extension is not loading even if I do |
| @@ -14,11 +14,18 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" | |||
| SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c" | |||
| StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" | |||
|
|
|||
| [weakdeps] | |||
| NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" | |||
There was a problem hiding this comment.
Need to also add it to extra for backwards compatibility with v1.8
|
LBroyden doesn't need any other dependency. Will add it in a different PR |
NNlib dependency is needed to handle the batched matrix multiply properly. Also needs SciML/DiffEqBase.jl#875