Use Type{T} self-types for rules on types on Julia 1.14

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.26.1"
+version = "1.26.2"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -12,7 +12,7 @@ BenchmarkTools = "1.5"
 Compat = "4.15"
 FiniteDifferences = "0.12.32"
 OffsetArrays = "1.14"
-StaticArrays = "1.9.2"
+StaticArrays = "1.9.8"
 julia = "1.10"
 
 [extensions]

docs/src/rule_author/rule_definition_tools.md

@@ -11,10 +11,10 @@ For example, instead of manually defining the `frule` and the `rrule` for string
 ```
 defines the following `frule` and `rrule` automatically
 ```julia
-function ChainRulesCore.frule(var"##_#1600", ::Core.Typeof(*), String::Any...; kwargs...)
+function ChainRulesCore.frule(var"##_#1600", ::typeof(*), String::Any...; kwargs...)
     return (*(String...; kwargs...), NoTangent())
 end
-function ChainRulesCore.rrule(::Core.Typeof(*), String::Any...; kwargs...)
+function ChainRulesCore.rrule(::typeof(*), String::Any...; kwargs...)
     return (*(String...; kwargs...), function var"*_pullback"(_)
         (ZeroTangent(), ntuple((_->NoTangent()), 0 + length(String))...)
     end)

docs/src/rule_author/writing_good_rules.md

@@ -108,33 +108,6 @@ Use `Type{<:Bar}` (with the `<:`) for non-concrete types, such that the `rrule`
 In particular, be careful not to use `typeof(Bar)` here.
 Because `typeof(Bar)` is `DataType`, using this to define an `rrule`/`frule` will define an `rrule`/`frule` for all constructors.
 
-You can check which to use with `Core.Typeof`:
-
-```julia-repl
-julia> function foo end
-foo (generic function with 0 methods)
-
-julia> typeof(foo)
-typeof(foo)
-
-julia> Core.Typeof(foob)
-typeof(foo)
-
-julia> typeof(Bar)
-DataType
-
-julia> Core.Typeof(Bar)
-Type{Bar}
-
-julia> abstract type AbstractT end
-
-julia> typeof(AbstractT)
-DataType
-
-julia> Core.Typeof(AbstractT)
-Type{AbstractT}
-```
-
 ### Functors (callable objects)
 
 In contrast to defining a rule for a constructor, it is possible to define rules for calling an instance of an object.

src/rule_definition_tools.jl

@@ -5,6 +5,9 @@ macro strip_linenos(expr)
     return esc(Base.remove_linenums!(expr))
 end
 
+# `Core.Typeof` before Julia 1.14 (now `Core.TypeEqOf`), see JuliaLang/julia#62001
+_selftype(@nospecialize(x)) = isa(x, Type) ? Type{x} : typeof(x)
+
 ############################################################################################
 ### @scalar_rule
 
@@ -183,7 +186,7 @@ function derivatives_given_output end
 function scalar_derivative_expr(__source__, f, setup_stmts, inputs, partials)
     return @strip_linenos quote
         function ChainRulesCore.derivatives_given_output(
-            $(esc(:Ω)), ::Core.Typeof($f), $(inputs...)
+            $(esc(:Ω)), ::$_selftype($f), $(inputs...)
         )
             $(__source__)
             $(setup_stmts...)
@@ -215,7 +218,7 @@ function scalar_frule_expr(__source__, f, call, setup_stmts, inputs, partials)
     return @strip_linenos quote
         # _ is the input derivative w.r.t. function internals. since we do not
         # allow closures/functors with @scalar_rule, it is always ignored
-        function ChainRulesCore.frule((_, $(Δs...)), ::Core.Typeof($f), $(inputs...))
+        function ChainRulesCore.frule((_, $(Δs...)), ::$_selftype($f), $(inputs...))
             $(__source__)
             $(esc(:Ω)) = $call
             $(setup_stmts...)
@@ -255,7 +258,7 @@ function scalar_rrule_expr(__source__, f, call, setup_stmts, inputs, partials)
     end
 
     return @strip_linenos quote
-        function ChainRulesCore.rrule(::Core.Typeof($f), $(inputs...))
+        function ChainRulesCore.rrule(::$_selftype($f), $(inputs...))
             $(__source__)
             $(esc(:Ω)) = $call
             $(setup_stmts...)
@@ -607,7 +610,7 @@ function _split_primal_name(primal_name)
     is_qualified = Meta.isexpr(primal_name, :(.))
     is_parameterized = Meta.isexpr(primal_name, :curly)
     if is_plain || is_qualified || is_parameterized
-        primal_name_sig = :(::$Core.Typeof($primal_name))
+        primal_name_sig = :(::$_selftype($primal_name))
         return primal_name_sig, primal_name
     elseif Meta.isexpr(primal_name, :(::))  # e.g. (::T)(x, y)
         _primal_name = gensym(Symbol(:instance_, primal_name.args[end]))

test/projection.jl

@@ -457,12 +457,11 @@ struct NoSuperType end
         @test 0 == @ballocated $pdiag(dx) setup = (dx = Diagonal(rand(10^3)))
 
         psymm = ProjectTo(Symmetric(rand(10^3, 10^3)))
-        allocs = @ballocated $psymm(dx) setup = (dx = Symmetric(rand(10^3, 10^3)))  # 64
-        if VERSION > v"1.13"
-            @test 0 == allocs
-        else
-            @test_broken 0 == allocs
-        end
+        # The Union{Symmetric{…,Matrix}, Symmetric{…,ReshapedArray}} return value is only
+        # unboxed since JuliaLang/julia#55045 (sret_union ABI for pointer-ful types)
+        @test 0 == @ballocated($psymm(dx), setup = (dx = Symmetric(rand(10^3, 10^3)))) broken = (
+            VERSION < v"1.14.0-DEV.1348"
+        )
     end
 
     @testset "#685" begin

test/rule_definition_tools.jl

@@ -29,6 +29,10 @@ struct NonDiffCounterExample
     x
 end
 
+struct NonDiffTypeExample
+    x
+end
+
 module NonDiffModuleExample
 nondiff_2_1(x, y) = fill(7.5, 100)[x + y]
 end
@@ -142,6 +146,15 @@ end
             # make sure that isn't happenning
             @test frule((ZeroTangent(), 1.2), NonDiffCounterExample, 2.0) === nothing
             @test rrule(NonDiffCounterExample, 2.0) === nothing
+
+            # https://github.com/JuliaDiff/ChainRulesCore.jl/issues/713
+            # a more specific hand-written `::Type{T}` rule must not be ambiguous
+            @non_differentiable NonDiffTypeExample(::Any)
+            function ChainRulesCore.rrule(::Type{NonDiffTypeExample}, x::Float64)
+                return NonDiffTypeExample(x), Δ -> (NoTangent(), 2Δ)
+            end
+            @test rrule(NonDiffTypeExample, 2.0)[2](1.2) == (NoTangent(), 2.4)
+            @test rrule(NonDiffTypeExample, 2)[2](1.2) == (NoTangent(), NoTangent())
         end
 
         @testset "Varargs" begin