Merge pull request #10 from JuliaAI/static-transform

Fix handling of Static transformers

Author: ablaom

src/attemptors.jl

@@ -101,7 +101,11 @@ function operations(fitted_machine, data...; throw=false, verbosity=1)
             push!(operations, "predict")
         end
         if :transform in methods
-            W = transform(fitted_machine, first(data))
+            W = if model isa Static
+                transform(fitted_machine, data...)
+            else
+                transform(fitted_machine, first(data))
+            end
             model isa Static || transform(fitted_machine, rows=test)
             model isa Static || transform(fitted_machine, rows=:)
             push!(operations, "transform")

test/attemptors.jl

@@ -37,4 +37,32 @@ MLJBase.load_path(::Type{<:DummyModel}) = "DummyPackage.Some.Where.Over.The.Rain
     @test outcome == "✓"
 end
 
+struct DummyStatic <: Static end
+MLJBase.transform(::DummyStatic, _, x, y) = hcat(x, y)
+MLJBase.package_name(::Type{<:DummyStatic}) = "DummyPackage"
+MLJBase.load_path(::Type{<:DummyStatic}) = "DummyPackage.Some.Thing.Different"
+
+struct SupervisedTransformer <: Deterministic end
+MLJBase.fit(::SupervisedTransformer, verbosity, X, y) = (42, nothing, nothing)
+MLJBase.predict(::SupervisedTransformer, _, Xnew) = fill(4.5, length(Xnew))
+MLJBase.transform(model::SupervisedTransformer, Θ, Xnew) =
+    predict(model, Θ, Xnew)
+MLJBase.package_name(::Type{<:SupervisedTransformer}) = "DummyPackage"
+MLJBase.load_path(::Type{<:SupervisedTransformer}) =
+    "DummyPackage.Some.Thing.Else"
+
+@testset "operations" begin
+    X = fill(1.2, 10)
+    y = X
+    mach = machine(SupervisedTransformer(), X, y) |> fit!
+    operations, outcome = MLJTestInterface.operations(mach, X, y, throw=true)
+    @test operations == "predict, transform"
+    @test outcome == "✓"
+
+    smach = machine(DummyStatic())
+    operations, outcome = MLJTestInterface.operations(smach, X, y)
+    @test operations == "transform"
+    @test outcome == "✓"
+end
+
 true