allow, and test, integer x

Author: Michael Abbott

src/rulesets/LinearAlgebra/norm.jl

@@ -108,8 +108,7 @@ end
 
 function _normp_back_x(x, p, y, Δy)
     c = real(Δy) / y
-    ∂x = similar(x)
-    broadcast!(∂x, x) do xi
+    ∂x = broadcast(x) do xi
         a = norm(xi)
         ∂xi = xi * ((a / y)^(p - 2) * c)
         return ifelse(isfinite(∂xi), ∂xi, zero(∂xi))
@@ -174,8 +173,7 @@ function rrule(::typeof(LinearAlgebra.norm1), x::AbstractArray)
 end
 
 function _norm1_back(x, y, Δy)
-    ∂x = similar(x)
-    ∂x .= sign.(x) .* real(Δy)
+    ∂x = sign.(x) .* real(Δy)
     return ∂x
 end
 function _norm1_back!(∂x, x, y, Δy)
@@ -207,8 +205,7 @@ function _norm2_forward(x, Δx, y)
     return ∂y
 end
 function _norm2_back(x, y, Δy)
-    ∂x = similar(x)
-    ∂x .= x .* (real(Δy) * pinv(y))
+    ∂x = x .* (real(Δy) * pinv(y))
     return ∂x
 end
 function _norm2_back!(∂x, x, y, Δy)

test/rulesets/LinearAlgebra/norm.jl

@@ -50,6 +50,12 @@
             # frule_test(fnorm, (xp, ẋ))
             rrule_test(fnorm, ȳ, (xp, x̄))
         end
+        T == Float64 && ndims(x) == 1 && @testset "Integer input" begin
+            x = [1,2,3]
+            _, int_back = rrule(fnorm, x)
+            _, float_back = rrule(fnorm, float(x))
+            @test unthunk(int_back(1.0)[2]) ≈ unthunk(float_back(1.0)[2])
+        end
     end
     @testset "norm(x::Array{$T,$(length(sz))})" for
         T in (Float64, ComplexF64),
@@ -127,6 +133,12 @@
         ȳ = rand_tangent(fnorm(x, p))
         @test extern(rrule(fnorm, zero(x), p)[2](ȳ)[2]) ≈ zero(x)
         @test rrule(fnorm, x, p)[2](Zero())[2] isa Zero
+        T == Float64 && sz == (3,) && @testset "Integer input, p=$p" begin
+            x = [1,2,3]
+            _, int_back = rrule(fnorm, x, p)
+            _, float_back = rrule(fnorm, float(x), p)
+            @test unthunk(unthunk(int_back(1.0)[2])) ≈ unthunk(unthunk(float_back(1.0)[2]))
+        end
     end
     @testset "norm($fdual(::Vector{$T}), p)" for
         T in (Float64, ComplexF64),