From f72e220d9144e25fab4e1706357250f19c6b0a29 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 27 Aug 2024 12:56:00 +0530 Subject: [PATCH] Specialize broadcasting more unary functions over a OneElement (#384) * Specialize broadcasting more unary functions over a OneElement * Add tests --- src/oneelement.jl | 12 ++++-------- test/runtests.jl | 4 +++- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/src/oneelement.jl b/src/oneelement.jl index 39687a2a..a98267a6 100644 --- a/src/oneelement.jl +++ b/src/oneelement.jl @@ -397,14 +397,10 @@ end # broadcast -function broadcasted(::DefaultArrayStyle{N}, ::typeof(conj), r::OneElement{<:Any,N}) where {N} - OneElement(conj(r.val), r.ind, axes(r)) -end -function broadcasted(::DefaultArrayStyle{N}, ::typeof(real), r::OneElement{<:Any,N}) where {N} - OneElement(real(r.val), r.ind, axes(r)) -end -function broadcasted(::DefaultArrayStyle{N}, ::typeof(imag), r::OneElement{<:Any,N}) where {N} - OneElement(imag(r.val), r.ind, axes(r)) +for f in (:abs, :abs2, :conj, :real, :imag) + @eval function broadcasted(::DefaultArrayStyle{N}, ::typeof($f), r::OneElement{<:Any,N}) where {N} + OneElement($f(r.val), r.ind, axes(r)) + end end function broadcasted(::DefaultArrayStyle{N}, ::typeof(^), r::OneElement{<:Any,N}, x::Number) where {N} OneElement(r.val^x, r.ind, axes(r)) diff --git a/test/runtests.jl b/test/runtests.jl index cd6555b5..ee2fb9a7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -2669,9 +2669,11 @@ end end @testset "broadcasting" begin - for v in (OneElement(2, 3, 4), OneElement(2im, (1,2), (3,4))) + for v in (OneElement(-2, 3, 4), OneElement(2im, (1,2), (3,4))) w = Array(v) n = 2 + @test abs.(v) == abs.(w) + @test abs2.(v) == abs2.(w) @test real.(v) == real.(w) @test imag.(v) == imag.(w) @test conj.(v) == conj.(w)