Adding support for Equal (#118)

FluxML · Feb 7, 2025 · 7b3c4f9 · 7b3c4f9
1 parent 6132a74
commit 7b3c4f9
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/src/load.jl b/src/load.jl
@@ -87,6 +87,10 @@ function load_node!(tape::Tape, ::OpConfig{:ONNX, :Where}, args::VarVec, attrs::
     return push_call!(tape, _where, args...)
 end
 
+function load_node!(tape::Tape, ::OpConfig{:ONNX, :Equal}, args::VarVec, attrs::AttrDict)
+    return push_call!(tape, _equal, args...)
+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))

diff --git a/src/ops.jl b/src/ops.jl
@@ -72,6 +72,10 @@ function _where(condition, x, y)
     return ifelse.(condition, x, y)
 end
 
+function _equal(x, y)
+    return x .== y
+end
+
 add(xs...) = .+(xs...)
 sub(xs...) = .-(xs...)
 _sin(x) = sin.(x)

diff --git a/src/save.jl b/src/save.jl
@@ -161,6 +161,11 @@ function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(_where)}, op::Umlaut
     push!(g.node, nd)
 end
 
+function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(_equal)}, op::Umlaut.Call)
+    nd = NodeProto("Equal", 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)],

diff --git a/test/saveload.jl b/test/saveload.jl
@@ -120,6 +120,12 @@ import ONNX: NodeProto, ValueInfoProto, AttributeProto, onnx_name
         ort_test(ONNX._where, condition, A, B)
     end
 
+    @testset "Equal" begin
+        A = rand(Bool, (1, 20))
+        B = rand(Bool, (1, 20))
+        ort_test(ONNX._equal, A, B)
+    end
+
     @testset "Gemm" begin
         A, B, C = (rand(3, 4), rand(3, 4), rand(3, 3))
         ort_test(ONNX.onnx_gemm, A, B')