Adding support for Transpose

src/load.jl

@@ -71,6 +71,10 @@ function load_node!(tape::Tape, ::OpConfig{:ONNX, :And}, args::VarVec, attrs::At
     return push_call!(tape, and, args...)
 end
 
+function load_node!(tape::Tape, ::OpConfig{:ONNX, :Transpose}, args::VarVec, attrs::AttrDict)
+    return push_call!(tape, _transpose, args...; attrs...)
+end
+
 function load_node!(tape::Tape, nd::NodeProto, backend::Symbol)
     args = [tape.c.name2var[name] for name in nd.input]
     attrs = convert(Dict{Symbol, Any}, Dict(nd.attribute))

src/ops.jl

@@ -49,6 +49,15 @@ function and(x, y)
     return x .& y
 end
 
+function _transpose(x; perm = nothing)
+    if (perm === nothing)
+        return permutedims(x, (2,1))
+    else
+        _perm = perm .+ 1
+        return permutedims(x, _perm)
+    end
+end
+
 add(xs...) = .+(xs...)
 sub(xs...) = .-(xs...)
 _sin(x) = sin.(x)

src/save.jl

@@ -141,6 +141,11 @@ function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(and)}, op::Umlaut.Ca
     push!(g.node, nd)
 end
 
+function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(_transpose)}, op::Umlaut.Call)
+    nd = NodeProto("Transpose", op, kwargs2dict(op))
+    push!(g.node, nd)
+end
+
 function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(*)}, op::Umlaut.Call)
     nd = NodeProto(
         input=[onnx_name(v) for v in reverse(op.args)],

test/saveload.jl

@@ -60,6 +60,32 @@ import ONNX: NodeProto, ValueInfoProto, AttributeProto, onnx_name
         ort_test(ONNX._acosh, A)
     end
 
+    @testset "Transpose" begin
+        # ort_test() checks for expected output of functions, errors on Transpose 
+        # because of array shape; manually testing!
+
+        # Testing 2 dimension transpose
+        args = (rand(1, 2),)
+        size1 = size(first(args))
+        tape = Tape(ONNXCtx())
+        inp = [push!(tape, Input(arg)) for arg in args]
+        res = push_call!(tape, ONNX._transpose, inp...)
+        tape.result = res
+        size2 = size(play!(tape, first(args)))
+        @test size2 == (2, 1)
+
+        # Testing 3+ dimension transpose
+        args = (rand(3, 4, 5),)
+        kwargs = (perm = (1, 2, 0),)
+        size1 = size(first(args))
+        tape = Tape(ONNXCtx())
+        inp = [push!(tape, Input(arg)) for arg in args]
+        res = push_call!(tape, ONNX._transpose, inp...; kwargs...)
+        tape.result = res
+        size2 = size(play!(tape, first(args)))
+        @test size2 == (4, 5, 3)
+    end
+
     @testset "Gemm" begin
         A, B, C = (rand(3, 4), rand(3, 4), rand(3, 3))
         ort_test(ONNX.onnx_gemm, A, B')