Let user specify Integer ->Float conversion type

.github/workflows/CI.yml

@@ -52,6 +52,9 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@latest
       - uses: julia-actions/julia-runtest@latest
+      - uses: julia-actions/julia-runtest@latest # Run for Int -> Float32
+        env:
+          INT2FLOAT: Float32
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v1
         with:

Project.toml

@@ -5,6 +5,7 @@ version = "1.16.0"
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
@@ -21,6 +22,7 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [targets]
-test = ["Test", "BenchmarkTools", "FiniteDifferences", "OffsetArrays", "StaticArrays"]
+test = ["Test", "BenchmarkTools", "FiniteDifferences", "OffsetArrays", "StaticArrays", "UUIDs"]

src/ChainRulesCore.jl

@@ -17,6 +17,7 @@ export ignore_derivatives, @ignore_derivatives
 # tangents
 export Tangent, NoTangent, InplaceableThunk, Thunk, ZeroTangent, AbstractZero, AbstractThunk
 
+include("preferences.jl")
 include("debug_mode.jl")
 
 include("tangent_types/abstract_tangent.jl")

src/preferences.jl

@@ -0,0 +1,15 @@
+using Preferences
+
+const int2float_type = @load_preference("int2float", "Float64")
+
+function set_int2float!(int2float_type::String)
+    if !(int2float_type in ("Float64", "Float32", "Float16"))
+        throw(ArgumentError("Invalid int2float type: \"$(int2float_type)\""))
+    end
+
+    # Set it in our runtime values, as well as saving it to disk
+    @set_preferences!("int2float" => int2float_type)
+    @info(
+        "New int2float type set; restart your Julia session for this change to take effect!"
+    )
+end

src/projection.jl

@@ -159,7 +159,16 @@ ProjectTo(::Real) = ProjectTo{Real}()
 ProjectTo(::Complex) = ProjectTo{Complex}()
 ProjectTo(::Number) = ProjectTo{Number}()
 
-ProjectTo(x::Integer) = ProjectTo(float(x))
+function ProjectTo(x::Integer)
+    @static if int2float_type == "Float16"
+        return ProjectTo(Float16(x))
+    elseif int2float_type == "Float32"
+        return ProjectTo(Float32(x))
+    else # Float64
+        return ProjectTo(Float64(x))
+    end
+end
+
 ProjectTo(x::Complex{<:Integer}) = ProjectTo(float(x))
 
 # Preserve low-precision floats as accidental promotion is a common performance bug
@@ -385,7 +394,6 @@ function (project::ProjectTo{<:Tangent{<:Tuple}})(dx::AbstractArray)
     end
 end
 
-
 #####
 ##### `LinearAlgebra`
 #####

test/projection.jl

@@ -35,6 +35,8 @@ struct NoSuperType end
         @test ProjectTo(1.0f0 + 2im)(3) === 3.0f0 + 0im
         @test ProjectTo(big(1.0))(2) === 2
         @test ProjectTo(1.0)(2) === 2.0
+        @test ProjectTo(1.0f0)(2) === 2.0f0
+        @test ProjectTo(1)(2.0f0) === int2float(2.0)
 
         # Tangents
         ProjectTo(1.0f0 + 2im)(Tangent{ComplexF64}(; re=1, im=NoTangent())) ===
@@ -50,7 +52,7 @@ struct NoSuperType end
         # real & complex
         @test ProjectTo(1.0 + 1im)(Dual(1.0, 2.0)) isa Complex{<:Dual}
         @test ProjectTo(1.0 + 1im)(Complex(Dual(1.0, 2.0), Dual(1.0, 2.0))) isa
-              Complex{<:Dual}
+            Complex{<:Dual}
         @test ProjectTo(1.0)(Complex(Dual(1.0, 2.0), Dual(1.0, 2.0))) isa Dual
 
         # Tangent
@@ -59,10 +61,15 @@ struct NoSuperType end
 
     @testset "Base: arrays of numbers" begin
         pvec3 = ProjectTo([1, 2, 3])
-        @test pvec3(1.0:3.0) === 1.0:3.0
+        if int2float(1) isa Float64
+            @test pvec3(1.0:3.0) === 1.0:3.0
+        else
+            @test pvec3(1.0:3.0) == int2float.(1.0:3.0)
+        end
+
         @test pvec3(1:3) == 1.0:3.0  # would prefer ===, map(Float64, dx) would do that, not important
         @test pvec3([1, 2, 3 + 4im]) == 1:3
-        @test eltype(pvec3([1, 2, 3.0f0])) === Float64
+        @test eltype(pvec3([1, 2, 3.0f0])) === typeof(int2float(1))
 
         # reshape
         @test pvec3(reshape([1, 2, 3], 3, 1)) isa Vector
@@ -104,7 +111,7 @@ struct NoSuperType end
         @test ProjectTo([(1, 2), (3, 4), (5, 6)]) isa ProjectTo{AbstractArray}
 
         @test ProjectTo(Any[1, 2])(1:2) == [1.0, 2.0]  # projects each number.
-        @test Tuple(ProjectTo(Any[1, 2 + 3im])(1:2)) === (1.0, 2.0 + 0.0im)
+        @test Tuple(ProjectTo(Any[1, 2 + 3im])(1:2)) === (int2float(1.0), 2.0 + 0.0im)
         @test ProjectTo(Any[true, false]) isa ProjectTo{NoTangent}
 
         # empty arrays
@@ -143,7 +150,7 @@ struct NoSuperType end
 
         @test ProjectTo(Ref(true)) isa ProjectTo{NoTangent}
         @test ProjectTo(Ref([false]')) isa ProjectTo{NoTangent}
-        
+
         @test ProjectTo(Ref(1.0))(Ref(NoTangent())) === NoTangent()  # collapse all-zero
     end
 
@@ -154,7 +161,7 @@ struct NoSuperType end
             @test @inferred(pt1(pt1((1,)))) == pt1(pt1((1,)))            # accepts correct Tangent
             @test @inferred(pt1(Tangent{Any}(1))) == pt1((1,))           # accepts Tangent{Any}
         end
-        @test pt1([1,]) == Tangent{Tuple{Float64}}(1.0,)  # accepts Vector
+        @test pt1([1]) == Tangent{Tuple{Float64}}(1.0)  # accepts Vector
         @test @inferred(pt1(NoTangent())) === NoTangent()
         @test @inferred(pt1(ZeroTangent())) === ZeroTangent()
         @test @inferred(pt1((NoTangent(),))) === NoTangent()  # collapse all-zero
@@ -163,7 +170,9 @@ struct NoSuperType end
         @test_throws Exception pt1([])
 
         pt3 = ProjectTo(([1, 2, 3], false, :gamma)) # partly non-differentiable
-        @test pt3((1:3, 4, 5)) == Tangent{Tuple{Vector{Int}, Bool, Symbol}}([1.0, 2.0, 3.0], NoTangent(), NoTangent())
+        @test pt3((1:3, 4, 5)) == Tangent{Tuple{Vector{Int},Bool,Symbol}}(
+            [1.0, 2.0, 3.0], NoTangent(), NoTangent()
+        )
         @test ProjectTo((true, [false])) isa ProjectTo{NoTangent}
     end
 
@@ -215,15 +224,15 @@ struct NoSuperType end
 
     @testset "UniformScaling" begin
         @test ProjectTo(I)(123) === NoTangent()
-        @test ProjectTo(2 * I)(I * 3im) === 0.0 * I
-        @test ProjectTo((4 + 5im) * I)(Tangent{typeof(im * I)}(; λ = 6)) === (6.0 + 0.0im) * I
+        @test ProjectTo(2 * I)(I * 3im) === int2float(0.0) * I
+        @test ProjectTo((4 + 5im) * I)(Tangent{typeof(im * I)}(; λ=6)) === (6.0 + 0.0im) * I
         @test ProjectTo(7 * I)(Tangent{typeof(2I)}()) == ZeroTangent()
     end
 
     @testset "LinearAlgebra: $adj vectors" for adj in [transpose, adjoint]
         # adjoint vectors
         padj = ProjectTo(adj([1, 2, 3]))
-        adjT = typeof(adj([1, 2, 3.0]))
+        adjT = typeof(adj(int2float.([1, 2, 3.0])))
         @test padj(transpose(1:3)) isa adjT
         @test padj([4 5 6 + 7im]) isa adjT
         @test padj([4.0 5.0 6.0]) isa adjT
@@ -293,8 +302,13 @@ struct NoSuperType end
         pdiag = ProjectTo(Diagonal(1:3))
         @test pdiag(reshape(1:9, 3, 3)) == Diagonal([1, 5, 9])
         @test pdiag(pdiag(reshape(1:9, 3, 3))) == pdiag(reshape(1:9, 3, 3))
-        @test pdiag(rand(ComplexF32, 3, 3)) isa Diagonal{Float64}
-        @test pdiag(Diagonal(1.0:3.0)) === Diagonal(1.0:3.0)
+        @test pdiag(rand(ComplexF32, 3, 3)) isa Diagonal{Float32}
+        if int2float(1) isa Float64
+            @test pdiag(Diagonal(1.0:3.0)) === Diagonal(int2float.(1.0:3.0))
+        else
+            @test pdiag(Diagonal(1.0:3.0)) == Diagonal(int2float.(1.0:3.0))
+        end
+
         @test ProjectTo(Diagonal(randn(3) .> 0))(randn(3, 3)) == NoTangent()
         @test ProjectTo(Diagonal(randn(3) .> 0))(Diagonal(rand(3))) == NoTangent()
 
@@ -375,29 +389,29 @@ struct NoSuperType end
         pvec3 = ProjectTo([1, 2, 3])
         @test axes(pvec3(OffsetArray(rand(3), 0:2))) == (1:3,)
         @test pvec3(OffsetArray(rand(3), 0:2)) isa Vector  # relies on axes === axes test 
-        @test pvec3(OffsetArray(rand(3,1), 0:2, 0:0)) isa Vector
+        @test pvec3(OffsetArray(rand(3, 1), 0:2, 0:0)) isa Vector
     end
 
     #####
     ##### `StaticArrays`
     #####
 
-        @testset "StaticArrays" begin
-            # There is no code for this, but when argument isa StaticArray, axes(x) === axes(dx)
-            # implies a check, and reshape will wrap a Vector into a static SizedVector:
-            pstat = ProjectTo(SA[1, 2, 3])
-            @test axes(pstat(rand(3))) === (SOneTo(3),)
-
-            # This recurses into structured arrays:
-            pst = ProjectTo(transpose(SA[1, 2, 3]))
-            @test axes(pst(rand(1,3))) === (SOneTo(1), SOneTo(3))
-            @test pst(rand(1,3)) isa Transpose
-
-            # When the argument is an ordinary Array, static gradients are allowed to pass,
-            # like FillArrays. Collecting to an Array would cost a copy.
-            pvec3 = ProjectTo([1, 2, 3])
-            @test pvec3(SA[1, 2, 3]) isa StaticArray
-        end
+    @testset "StaticArrays" begin
+        # There is no code for this, but when argument isa StaticArray, axes(x) === axes(dx)
+        # implies a check, and reshape will wrap a Vector into a static SizedVector:
+        pstat = ProjectTo(SA[1, 2, 3])
+        @test axes(pstat(rand(3))) === (SOneTo(3),)
+
+        # This recurses into structured arrays:
+        pst = ProjectTo(transpose(SA[1, 2, 3]))
+        @test axes(pst(rand(1, 3))) === (SOneTo(1), SOneTo(3))
+        @test pst(rand(1, 3)) isa Transpose
+
+        # When the argument is an ordinary Array, static gradients are allowed to pass,
+        # like FillArrays. Collecting to an Array would cost a copy.
+        pvec3 = ProjectTo([1, 2, 3])
+        @test pvec3(SA[1, 2, 3]) isa StaticArray
+    end
 
     #####
     ##### `ChainRulesCore`
@@ -453,10 +467,12 @@ struct NoSuperType end
     end
 
     @testset "display" begin
+        int2float_type_ = typeof(int2float(2))
         @test repr(ProjectTo(1.1)) == "ProjectTo{Float64}()"
         @test occursin("ProjectTo{AbstractArray}(element", repr(ProjectTo([1, 2, 3])))
         str = repr(ProjectTo([1, 2, 3]'))
-        @test eval(Meta.parse(str))(ones(1, 3)) isa Adjoint{Float64,Vector{Float64}}
+        @test eval(Meta.parse(str))(ones(1, 3)) isa
+            Adjoint{int2float_type_,Vector{int2float_type_}}
     end
 
     VERSION > v"1.1" && @testset "allocation tests" begin

test/runtests.jl

@@ -1,12 +1,30 @@
 using Base.Broadcast: broadcastable
 using BenchmarkTools
+using Preferences
+using UUIDs
+
+# Test Float32 value for int2float
+
+if "INT2FLOAT" ∈ keys(ENV)
+    env_int2float = ENV["INT2FLOAT"]
+
+    if env_int2float ∈ ["Float32", "Float16"]
+        chainrulescore_uuid = UUID("d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4")
+        set_preferences!(chainrulescore_uuid, "int2float" => env_int2float)
+        println("")
+        println("Running ChainRulesCore tests with $env_int2float")
+    end
+end
+
 using ChainRulesCore
 using LinearAlgebra
 using LinearAlgebra.BLAS: ger!, gemv!, gemv, scal!
 using StaticArrays
 using SparseArrays
 using Test
 
+int2float(x) = ProjectTo(1)(x)
+
 @testset "ChainRulesCore" begin
     @testset "differentials" begin
         include("tangent_types/abstract_zero.jl")

test/tangent_types/notimplemented.jl

@@ -104,7 +104,7 @@
         f̄, x̄1, x̄2 = pb(3.1)
         @test f̄ == NoTangent()
         @test x̄1 isa ChainRulesCore.NotImplemented
-        @test x̄2 == 3.1
+        @test x̄2 == int2float(3.1)
 
         notimplemented2(x, y) = (x + y, x - y)
         @scalar_rule notimplemented2(x, y) (@not_implemented("notimplemented2"), 1) (1, -1)
@@ -119,6 +119,6 @@
         f̄, x̄1, x̄2 = pb((3.1, 4.5))
         @test f̄ == NoTangent()
         @test x̄1 isa ChainRulesCore.NotImplemented
-        @test x̄2 == -1.4
+        @test x̄2 == int2float(-1.4)
     end
 end