Adding support for ConstantOfShape (#121)

Author: dstarkenburg

src/load.jl

@@ -95,6 +95,11 @@ function load_node!(tape::Tape, ::OpConfig{:ONNX, :Neg}, args::VarVec, attrs::At
     return push_call!(tape, neg, args[1])
 end
 
+function load_node!(tape::Tape, ::OpConfig{:ONNX, :ConstantOfShape}, args::VarVec, attrs::AttrDict)
+    return push_call!(tape, makeshape, 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

@@ -76,6 +76,10 @@ function _equal(x, y)
     return x .== y
 end
 
+function makeshape(x; value = zeros(Float32, 1))
+    return fill(value..., x...)
+end
+
 add(xs...) = .+(xs...)
 sub(xs...) = .-(xs...)
 neg(x) = .-(x)

src/save.jl

@@ -171,6 +171,11 @@ function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(neg)}, op::Umlaut.Ca
     push!(g.node, nd)
 end
 
+function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(makeshape)}, op::Umlaut.Call)
+    nd = NodeProto("ConstantOfShape", 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

@@ -131,6 +131,28 @@ import ONNX: NodeProto, ValueInfoProto, AttributeProto, onnx_name
         ort_test(ONNX.neg, A)
     end
 
+    @testset "ConstantOfShape" begin
+        # ort_test() checks for expected output of functions, errors on ConstantOfShape 
+        # because of array shape; manually testing!
+
+        # Testing expansion of shape in ConstantOfShape
+        args = [2, 3]
+        attrs = randn(Float32, 1)
+        tape = Tape(ONNXCtx())
+        inp = push!(tape, Input(args))
+        res = push_call!(tape, ONNX.makeshape, inp; value = attrs)
+        tape.result = res
+        
+        # Make sure size is desired shape
+        @test size(play!(tape, args)) == (2, 3)
+        
+        # Make sure elements are of the desired datatype
+        @test eltype(attrs) == eltype(play!(tape, args))
+
+        # Make sure the output is filled with the correct value
+        @test attrs[1] == play!(tape, args)[1]
+    end
+
     @testset "Gemm" begin
         A, B, C = (rand(3, 4), rand(3, 4), rand(3, 3))
         ort_test(ONNX.onnx_gemm, A, B')