add sol_kwargs to mlj models

impICNF · Oct 15, 2024 · dc15bc4 · dc15bc4
1 parent 315881d
commit dc15bc4
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 8 deletions.
diff --git a/Project.toml b/Project.toml
@@ -69,4 +69,4 @@ SciMLSensitivity = "7"
 ScientificTypesBase = "3"
 Statistics = "1"
 Zygote = "0.6"
-julia = "1.9"
+julia = "1.10"
diff --git a/benchmark/Project.toml b/benchmark/Project.toml
@@ -19,4 +19,4 @@ Lux = "1"
 PkgBenchmark = "0.2"
 StableRNGs = "1"
 Zygote = "0.6"
-julia = "1.9"
+julia = "1.10"
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -3,4 +3,4 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 
 [compat]
 Documenter = "1"
-julia = "1.9"
+julia = "1.10"
diff --git a/src/exts/mlj_ext/core_cond_icnf.jl b/src/exts/mlj_ext/core_cond_icnf.jl
@@ -8,6 +8,7 @@ mutable struct CondICNFModel{AICNF <: AbstractICNF} <: MLJICNF{AICNF}
 
     use_batch::Bool
     batch_size::Int
+    sol_kwargs::NamedTuple
 end
 
 function CondICNFModel(
@@ -18,8 +19,9 @@ function CondICNFModel(
     adtype::ADTypes.AbstractADType = ADTypes.AutoZygote(),
     use_batch::Bool = true,
     batch_size::Int = 32,
+    sol_kwargs::NamedTuple = (;),
 )
-    CondICNFModel(m, loss, optimizers, n_epochs, adtype, use_batch, batch_size)
+    CondICNFModel(m, loss, optimizers, n_epochs, adtype, use_batch, batch_size, sol_kwargs)
 end
 
 function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
@@ -64,8 +66,8 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
         tst_epochs = @timed res = SciMLBase.solve(
             optprob_re,
             opt;
-            progress = true,
             epochs = model.n_epochs,
+            model.sol_kwargs...,
         )
         ps .= res.u
         @info(

diff --git a/src/exts/mlj_ext/core_icnf.jl b/src/exts/mlj_ext/core_icnf.jl
@@ -8,6 +8,7 @@ mutable struct ICNFModel{AICNF <: AbstractICNF} <: MLJICNF{AICNF}
 
     use_batch::Bool
     batch_size::Int
+    sol_kwargs::NamedTuple
 end
 
 function ICNFModel(
@@ -18,8 +19,9 @@ function ICNFModel(
     adtype::ADTypes.AbstractADType = ADTypes.AutoZygote(),
     use_batch::Bool = true,
     batch_size::Int = 32,
+    sol_kwargs::NamedTuple = (;),
 )
-    ICNFModel(m, loss, optimizers, n_epochs, adtype, use_batch, batch_size)
+    ICNFModel(m, loss, optimizers, n_epochs, adtype, use_batch, batch_size, sol_kwargs)
 end
 
 function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
@@ -62,8 +64,8 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
         tst_epochs = @timed res = SciMLBase.solve(
             optprob_re,
             opt;
-            progress = true,
             epochs = model.n_epochs,
+            model.sol_kwargs...,
         )
         ps .= res.u
         @info(

diff --git a/test/Project.toml b/test/Project.toml
@@ -43,4 +43,4 @@ StableRNGs = "1"
 TerminalLoggers = "0.1"
 Zygote = "0.6"
 cuDNN = "1"
-julia = "1.9"
+julia = "1.10"