Specialize broadcasting more unary functions over a OneElement (#384)

* Specialize broadcasting more unary functions over a OneElement * Add tests
JuliaArrays · Aug 27, 2024 · f72e220 · f72e220
1 parent 414dba7
commit f72e220
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 9 deletions.
diff --git 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
@@ -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)