Skip to content

Commit

Permalink
Adding support for Cos (#108)
Browse files Browse the repository at this point in the history
* Added support for Cos

* Adding test case for Cos
  • Loading branch information
dstarkenburg authored Dec 6, 2024
1 parent ed396c2 commit b5b9f52
Show file tree
Hide file tree
Showing 4 changed files with 16 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/load.jl
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,10 @@ function load_node!(tape::Tape, ::OpConfig{:ONNX, :Sin}, args::VarVec, attrs::At
return push_call!(tape, _sin, args[1])
end

function load_node!(tape::Tape, ::OpConfig{:ONNX, :Cos}, args::VarVec, attrs::AttrDict)
return push_call!(tape, _cos, args[1])
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))
Expand Down
1 change: 1 addition & 0 deletions src/ops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ end
add(xs...) = .+(xs...)
sub(xs...) = .-(xs...)
_sin(x) = sin.(x)
_cos(x) = cos.(x)
mul(xs...) = .*(xs...)
relu(x) = NNlib.relu.(x)
leakyrelu(x;a = 0.01) = NNlib.leakyrelu.(x,a)
Expand Down
5 changes: 5 additions & 0 deletions src/save.jl
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,11 @@ function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(_sin)}, op::Umlaut.C
push!(g.node, nd)
end

function save_node!(g::GraphProto, ::OpConfig{:ONNX, typeof(_cos)}, op::Umlaut.Call)
nd = NodeProto("Cos", 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)],
Expand Down
6 changes: 6 additions & 0 deletions test/saveload.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,12 @@ import ONNX: NodeProto, ValueInfoProto, AttributeProto, onnx_name
ort_test(ONNX._sin, A)
end

@testset "Cos" begin
# ONNXRunTime has no implementation for Cos(x::Float64), using Float32
A = rand(Float32, 3, 4)
ort_test(ONNX._cos, A)
end

@testset "Gemm" begin
A, B, C = (rand(3, 4), rand(3, 4), rand(3, 3))
ort_test(ONNX.onnx_gemm, A, B')
Expand Down

0 comments on commit b5b9f52

Please sign in to comment.