Fix split, uncomment tests (#98)

FluxML · Oct 22, 2023 · 7e8d640 · 7e8d640
1 parent fb7e66d
commit 7e8d640
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 19 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ONNX"
 uuid = "d0dd6a25-fac6-55c0-abf7-829e0c774d20"
-version = "0.2.5"
+version = "0.2.6"
 
 [deps]
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"

diff --git a/src/load.jl b/src/load.jl
@@ -222,21 +222,21 @@ function load_node!(tape::Tape, ::OpConfig{:ONNX, :Slice}, args::VarVec, attrs::
     return push_call!(tape, onnx_slice, args...)
 end
 
-function load_node!(tape::Tape, ::OpConfig{:ONNX, :Split}, inputs::VarVec, attrs::AttrDict)
+function load_node!(tape::Tape, ::OpConfig{:ONNX, :Split}, args::VarVec, attrs::AttrDict)
     axis = get(attrs, :axis, 0)
     split = if haskey(attrs, :split) # Version 1, 2, 11
         attrs[:split]
     elseif length(args) == 2
-        inputs[2]
+        args[2]
     else
         # the results cannot be split in multiple outputs on the tape
         # if the output size is not known during tracing.
         error("Unhandled case where split is not provided")
     end
-    out = push_call!(tape, onnx_split, first(inputs), split; axis)
+    out = push_call!(tape, onnx_split, first(args), split; axis)
     return Tuple(
         push_call!(tape, getfield, out, i)
-        for i in eachindex(split)
+        for i in eachindex(split.op.val)
     )
 end
 

diff --git a/test/saveload.jl b/test/saveload.jl
@@ -186,18 +186,17 @@ import ONNX: NodeProto, ValueInfoProto, AttributeProto, onnx_name
         ort_test(ONNX.onnx_concat, [1 2 3; 1 2 3], [4  5; 4 5]; axis=0)
     end
 
-    # TODO: Split is not implemented in ONNXRuntime.jl
-    # @testset "Split" begin
-    #     x = rand(3, 20, 10); split = [5, 10, 5];
-    #     args = (x, split)
-    #     tape = Tape(ONNXCtx())
-    #     inp = [push!(tape, Input(a)) for a in args]
-    #     out = push_call!(tape, ONNX.onnx_split, inp...; axis=1)
-    #     push_call!(tape, getfield, out, 1)
-    #     push_call!(tape, getfield, out, 2)
-    #     push_call!(tape, getfield, out, 3)
-    #     tape.result = out
-
-    #     ort_test(tape, args...)
-    # end
+    @testset "Split" begin
+        x = rand(3, 20, 10); split = [5, 10, 5];
+        args = (x, split)
+        tape = Tape(ONNXCtx())
+        inp = [push!(tape, Input(a)) for a in args]
+        out = push_call!(tape, ONNX.onnx_split, inp...; axis=1)
+        push_call!(tape, getfield, out, 1)
+        push_call!(tape, getfield, out, 2)
+        push_call!(tape, getfield, out, 3)
+        tape.result = out
+
+        ort_test(tape, args...)
+    end
 end