fix type check

impICNF · May 23, 2024 · 0ec4c19 · 0ec4c19
1 parent 9705a13
commit 0ec4c19
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 10 deletions.
diff --git a/src/exts/mlj_ext/core_cond_icnf.jl b/src/exts/mlj_ext/core_cond_icnf.jl
@@ -43,7 +43,7 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
     tst_overall = @timed for opt in model.optimizers
         tst_epochs = @timed for ep in 1:(model.n_epochs)
             if model.use_batch
-                if model.m.compute_mode <: VectorMode
+                if model.m.compute_mode isa VectorMode
                     data = MLUtils.DataLoader(
                         (x, y);
                         batchsize = -1,
@@ -52,7 +52,7 @@ function MLJModelInterface.fit(model::CondICNFModel, verbosity, XY)
                         parallel = false,
                         buffer = false,
                     )
-                elseif model.m.compute_mode <: MatrixMode
+                elseif model.m.compute_mode isa MatrixMode
                     data = MLUtils.DataLoader(
                         (x, y);
                         batchsize = model.batch_size,
@@ -110,13 +110,13 @@ function MLJModelInterface.transform(model::CondICNFModel, fitresult, XYnew)
     end
     (ps, st) = fitresult
 
-    tst = @timed if model.m.compute_mode <: VectorMode
+    tst = @timed if model.m.compute_mode isa VectorMode
         logp̂x = broadcast(
             (x, y) -> first(inference(model.m, TestMode(), x, y, ps, st)),
             eachcol(xnew),
             eachcol(ynew),
         )
-    elseif model.m.compute_mode <: MatrixMode
+    elseif model.m.compute_mode isa MatrixMode
         logp̂x = first(inference(model.m, TestMode(), xnew, ynew, ps, st))
     else
         error("Not Implemented")

diff --git a/src/exts/mlj_ext/core_icnf.jl b/src/exts/mlj_ext/core_icnf.jl
@@ -41,7 +41,7 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
     tst_overall = @timed for opt in model.optimizers
         tst_epochs = @timed for ep in 1:(model.n_epochs)
             if model.use_batch
-                if model.m.compute_mode <: VectorMode
+                if model.m.compute_mode isa VectorMode
                     data = MLUtils.DataLoader(
                         (x,);
                         batchsize = -1,
@@ -50,7 +50,7 @@ function MLJModelInterface.fit(model::ICNFModel, verbosity, X)
                         parallel = false,
                         buffer = false,
                     )
-                elseif model.m.compute_mode <: MatrixMode
+                elseif model.m.compute_mode isa MatrixMode
                     data = MLUtils.DataLoader(
                         (x,);
                         batchsize = model.batch_size,
@@ -105,9 +105,9 @@ function MLJModelInterface.transform(model::ICNFModel, fitresult, Xnew)
     end
     (ps, st) = fitresult
 
-    tst = @timed if model.m.compute_mode <: VectorMode
+    tst = @timed if model.m.compute_mode isa VectorMode
         logp̂x = broadcast(x -> first(inference(model.m, TestMode(), x, ps, st)), eachcol(xnew))
-    elseif model.m.compute_mode <: MatrixMode
+    elseif model.m.compute_mode isa MatrixMode
         logp̂x = first(inference(model.m, TestMode(), xnew, ps, st))
     else
         error("Not Implemented")

diff --git a/test/call_tests.jl b/test/call_tests.jl
@@ -73,10 +73,10 @@ Test.@testset "Call Tests" begin
             Distributions.Beta{data_type}(convert(Tuple{data_type, data_type}, (2, 4))...)
         data_dist2 =
             Distributions.Beta{data_type}(convert(Tuple{data_type, data_type}, (4, 2))...)
-        if compute_mode <: ContinuousNormalizingFlows.VectorMode
+        if compute_mode isa ContinuousNormalizingFlows.VectorMode
             r = convert.(data_type, rand(data_dist, nvars))
             r2 = convert.(data_type, rand(data_dist2, nvars))
-        elseif compute_mode <: ContinuousNormalizingFlows.MatrixMode
+        elseif compute_mode isa ContinuousNormalizingFlows.MatrixMode
             r = convert.(data_type, rand(data_dist, nvars, ndata))
             r2 = convert.(data_type, rand(data_dist2, nvars, ndata))
         end