Merge pull request #395 from mcabbott/scalarmacro

Always include `ProjectTo` in `@scalar_rule`

Author: oxinabox

src/rule_definition_tools.jl

@@ -108,7 +108,7 @@ end
 
 returns (in order) the correctly escaped:
     - `call` with out any type constraints
-    - `setup_stmts`: the content of `@setup` or `nothing` if that is not provided,
+    - `setup_stmts`: the content of `@setup` or `[]` if that is not provided,
     -  `inputs`: with all args having the constraints removed from call, or
         defaulting to `Number`
     - `partials`: which are all `Expr{:tuple,...}`
@@ -118,9 +118,9 @@ function _normalize_scalarrules_macro_input(call, maybe_setup, partials)
     # Setup: normalizing input form etc
 
     if Meta.isexpr(maybe_setup, :macrocall) && maybe_setup.args[1] == Symbol("@setup")
-        setup_stmts = map(esc, maybe_setup.args[3:end])
+        setup_stmts = Any[esc(ex) for ex in maybe_setup.args[3:end]]
     else
-        setup_stmts = (nothing,)
+        setup_stmts = []
         partials = (maybe_setup, partials...)
     end
     @assert Meta.isexpr(call, :call)
@@ -185,10 +185,14 @@ function scalar_rrule_expr(__source__, f, call, setup_stmts, inputs, partials)
     # because this is a pull-back there is one per output of function
     Δs = _propagator_inputs(n_outputs)
 
+    # Make a projector for each argument
+    projs, psetup  = _make_projectors(call.args[2:end])
+    append!(setup_stmts, psetup)
+
     # 1 partial derivative per input
     pullback_returns = map(1:n_inputs) do input_i
         ∂s = [partial.args[input_i] for partial in partials]
-        propagation_expr(Δs, ∂s, true)
+        propagation_expr(Δs, ∂s, true, projs[input_i])
     end
 
     # Multi-output functions have pullbacks with a tuple input that will be destructured
@@ -215,14 +219,23 @@ end
 "Declares properly hygenic inputs for propagation expressions"
 _propagator_inputs(n) = [esc(gensym(Symbol(:Δ, i))) for i in 1:n]
 
+"given the variable names, escaped but without types, makes setup expressions for projection operators"
+function _make_projectors(xs)
+    projs = map(x -> Symbol(:proj_, x.args[1]), xs)
+    setups = map((x,p) -> :($p = ProjectTo($x)), xs, projs)
+    return projs, setups
+end
+
 """
-    propagation_expr(Δs, ∂s, _conj = false)
+    propagation_expr(Δs, ∂s, [_conj=false, proj=identity])
 
-    Returns the expression for the propagation of
-    the input gradient `Δs` though the partials `∂s`.
-    Specify `_conj = true` to conjugate the partials.
+Returns the expression for the propagation of
+the input gradient `Δs` though the partials `∂s`.
+Specify `_conj = true` to conjugate the partials.
+Projector `proj` is a function that will be applied at the end; 
+    for `rrules` it is usually a `ProjectTo(x)`, for `frules` it is `identity`
 """
-function propagation_expr(Δs, ∂s, _conj = false)
+function propagation_expr(Δs, ∂s, _conj=false, proj=identity)
     # This is basically Δs ⋅ ∂s
     _∂s = map(∂s) do ∂s_i
         if _conj
@@ -249,7 +262,7 @@ function propagation_expr(Δs, ∂s, _conj = false)
         :($(_∂s[1]) * $(Δs[1]))
     end
 
-    return summed_∂_mul_Δs
+    return :($proj($summed_∂_mul_Δs))
 end
 
 """

test/rule_definition_tools.jl

@@ -237,6 +237,18 @@ end
             @test ẏ isa Tangent{Tuple{Irrational{:π}, Float64}, Tuple{Float32, Float32}}
         end
 
+        @testset "@scalar_rule projection" begin
+            make_imaginary(x) = im*x
+            @scalar_rule make_imaginary(x) im
+
+            # note: the === will make sure that these are Float64, not ComplexF64
+            @test (NoTangent(), 1.0) === rrule(make_imaginary, 2.0)[2](1.0*im)
+            @test (NoTangent(), 0.0) === rrule(make_imaginary, 2.0)[2](1.0)
+
+            @test (NoTangent(), 1.0+0.0im) === rrule(make_imaginary, 2.0im)[2](1.0*im)
+            @test (NoTangent(), 0.0-1.0im) === rrule(make_imaginary, 2.0im)[2](1.0)
+        end
+
         @testset "Regression tests against #276 and #265" begin
             # https://github.com/JuliaDiff/ChainRulesCore.jl/pull/276
             # https://github.com/JuliaDiff/ChainRulesCore.jl/pull/265