Merge pull request #847 from sethaxen/rmnorm

Remove adjoint for norm

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.5.15"
+version = "0.5.16"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -26,7 +26,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 AbstractFFTs = "0.5"
 ArrayLayouts = "0.1, 0.2, 0.3, 0.4"
-ChainRules = "0.7.33"
+ChainRules = "0.7.34"
 DiffRules = "1.0"
 FillArrays = "0.8, 0.9, 0.10"
 ForwardDiff = "0.10"

src/lib/array.jl

@@ -421,11 +421,6 @@ end
   end
 end
 
-function _pullback(cx::AContext, ::typeof(norm), x::AbstractArray, p::Real = 2)
-  fallback = (x, p) -> sum(abs.(x).^p .+ eps(0f0)) ^ (one(eltype(x)) / p) # avoid d(sqrt(x))/dx == Inf at 0
-  _pullback(cx, fallback, x, p)
-end
-
 # LinAlg Matrix Types
 # ===================
 

test/gradcheck.jl

@@ -1643,4 +1643,21 @@ end
   end
 end
 
-@test gradient(x -> norm(x), rand(Float32, 2, 2))[1] isa Matrix{Float32}
+@testset "norm" begin
+    # rrule for norm is defined in ChainRules. These tests just check various norm-related
+    # issues are resolved
+
+    # check that type is not unnecessarily promoted
+    # https://github.com/FluxML/Zygote.jl/issues/663
+    @test gradient(norm, randn(Float32, 2, 2)) isa Tuple{Matrix{Float32}}
+    @test gradient(norm, randn(Float32, 2, 2), 3) isa Tuple{Matrix{Float32},Float32}
+    @test gradient(norm, randn(Float32, 2, 2), 3f0) isa Tuple{Matrix{Float32},Float32}
+    @test gradient(norm, randn(ComplexF32, 2, 2), 3.5f0) isa Tuple{Matrix{ComplexF32},Float32}
+
+    # just check that these do not error
+    # https://github.com/FluxML/Zygote.jl/issues/331
+    gradient(x->norm(x*[1, 1]), 1.23)
+    gradient(x->norm(x*[1 1]), 1.23)
+    gradient(x->norm(x*[1im, 1]), 1.23)
+    gradient(x->norm(x*[1im 1]), 1.23)
+end