Merge pull request #1006 from FluxML/ox/typeonlyrrules

use rrules even when all the arguments are types

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