Merge pull request #1006 from FluxML/ox/typeonlyrrules

use rrules even when all the arguments are types

Author: oxinabox

src/compiler/interface2.jl

@@ -7,9 +7,7 @@ function edge!(m::IRTools.Meta, edge::Core.MethodInstance)
 end
 
 @generated function _pullback(ctx::AContext, f, args...)
-  T = Tuple{f,args...}
-  ignore_sig(T) && return :(f(args...), Pullback{$T}(()))
-
+  # Try using ChainRulesCore
   if is_kwfunc(f, args...)
     # if it is_kw then `args[1]` are the keyword args, `args[2]` is actual function
     cr_T = Tuple{ZygoteRuleConfig{ctx}, args[2:end]...}
@@ -20,9 +18,12 @@ end
   end
 
   hascr, cr_edge = has_chain_rrule(cr_T)
-
   hascr && return :($chain_rrule_f(ZygoteRuleConfig(ctx), f, args...))
 
+  # No ChainRule, going to have to work it out.
+  T = Tuple{f,args...}
+  ignore_sig(T) && return :(f(args...), Pullback{$T}(()))
+
   g = try _generate_pullback_via_decomposition(T) catch e e end
   g === nothing && return :(f(args...), Pullback{$T}((f,)))
   meta, forw, _ = g

test/chainrules.jl

@@ -214,6 +214,24 @@ using ChainRulesCore, ChainRulesTestUtils, Zygote
         @test (nothing,) == Zygote.gradient(x->not_diff_kw_eg(x, 2), 10.4)
         @test (nothing,) == Zygote.gradient(x->not_diff_kw_eg(x, 2; kw=2.0), 10.4)
     end
+
+    @testset "Type only rrule" begin
+        struct StructForTestingTypeOnlyRRules{T}
+            x::T
+        end
+        StructForTestingTypeOnlyRRules() = StructForTestingTypeOnlyRRules(1.0)
+        
+        function ChainRulesCore.rrule(P::Type{<:StructForTestingTypeOnlyRRules})
+            # notice here we mess with the primal doing 2.0 rather than 1.0, this is for testing purposes
+            # and also because apparently people actually want to do this. Weird, but 🤷
+            # https://github.com/SciML/SciMLBase.jl/issues/69#issuecomment-865639754
+            P(2.0),  _ -> (NoTangent(),)
+        end
+
+        @assert StructForTestingTypeOnlyRRules().x == 1.0
+        aug_primal_val, _ = Zygote.pullback(x->StructForTestingTypeOnlyRRules(), 1.2)
+        @test aug_primal_val.x == 2.0
+    end
 end
 
 @testset "ChainRulesCore.rrule_via_ad" begin