Interpolate more in rule helpers and fix escaping of `@non_differentiable` #325

devmotion · 2021-03-31T10:41:17Z

The PR interpolates ~~all types and functions that are not provided by the user~~ more types and functions in the AST in the macros @scalar_rule and @non_differentiable. This fixes the following bug:

julia> using ChainRulesCore: @non_differentiable

julia> @non_differentiable length(x::AbstractVector)
ERROR: UndefVarError: ChainRulesCore not defined
Stacktrace:
 [1] top-level scope
   @ ~/.julia/packages/ChainRulesCore/ASgvC/src/rule_definition_tools.jl:327

Additionally, not the whole output of @non_differentiable is escaped anymore, fixing a possible name collision for arguments of name kwargs:

julia> using ChainRulesCore: ChainRulesCore, @non_differentiable

julia> @non_differentiable length(kwargs::AbstractVector)
ERROR: syntax: function argument name not unique: "kwargs" around /home/david/.julia/packages/ChainRulesCore/ASgvC/src/rule_definition_tools.jl:327
Stacktrace:
 [1] top-level scope
   @ REPL[17]:1

Macro output on the master branch:

julia> Base.remove_linenums!(@macroexpand @scalar_rule sincos(x) cos(x) -sin(x))
quote
    if !(sincos isa ChainRulesCore.Type) && ChainRulesCore.fieldcount(ChainRulesCore.typeof(sincos)) > 0
        ChainRulesCore.throw(ChainRulesCore.ArgumentError("@scalar_rule cannot be used on closures/functors (such as $(sincos))"))
    end
    begin
        function (ChainRulesCore.ChainRulesCore).frule((ChainRulesCore._, var"##Δ1#257"), ::ChainRulesCore.typeof(sincos), x::Number)
            Ω = sincos(x)
            nothing
            return (Ω, (ChainRulesCore.ChainRulesCore).Composite{ChainRulesCore.typeof(Ω)}(cos(x) * var"##Δ1#257", -(sin(x)) * var"##Δ1#257"))
        end
    end
    begin
        function (ChainRulesCore.ChainRulesCore).rrule(::ChainRulesCore.typeof(sincos), x::Number)
            Ω = sincos(x)
            nothing
            return (Ω, begin
                        function sincos_pullback((var"##Δ1#258", var"##Δ2#259"))
                            $(Expr(:meta, :inline))
                            return (ChainRulesCore.NO_FIELDS, ChainRulesCore.muladd.(ChainRulesCore.conj(-(sin(x))), var"##Δ2#259", ChainRulesCore.:*.(ChainRulesCore.conj(cos(x)), var"##Δ1#258")))
                        end
                    end)
        end
    end
end

julia> Base.remove_linenums!(@macroexpand @non_differentiable length(x::Vector))
quote
    function ChainRulesCore.frule(var"##_#260", ::Core.Typeof(length), x::Vector; kwargs...)
        return (length(x; kwargs...), DoesNotExist())
    end
    begin
        function (::Core.kwftype(typeof(ChainRulesCore.rrule)))(kwargs::Any, rrule::typeof(ChainRulesCore.rrule), ::Core.Typeof(length), x::Vector)
            return (length(x; kwargs...), function length_pullback(_)
                        (ChainRulesCore.DoesNotExist(), (ChainRulesCore.DoesNotExist(),)...)
                    end)
        end
        function ChainRulesCore.rrule(::Core.Typeof(length), x::Vector)
            return (length(x), function length_pullback(_)
                        (ChainRulesCore.DoesNotExist(), (ChainRulesCore.DoesNotExist(),)...)
                    end)
        end
    end
end

julia> Base.remove_linenums!(@macroexpand @non_differentiable length(x::Vector...))
quote
    function ChainRulesCore.frule(var"##_#261", ::Core.Typeof(length), x::Vector...; kwargs...)
        return (length(x...; kwargs...), DoesNotExist())
    end
    begin
        function (::Core.kwftype(typeof(ChainRulesCore.rrule)))(kwargs::Any, rrule::typeof(ChainRulesCore.rrule), ::Core.Typeof(length), x::Vector...)
            return (length(x...; kwargs...), function length_pullback(_)
                        (ChainRulesCore.DoesNotExist(), ntuple((_->ChainRulesCore.DoesNotExist()), 0 + length(x))...)
                    end)
        end
        function ChainRulesCore.rrule(::Core.Typeof(length), x::Vector...)
            return (length(x...), function length_pullback(_)
                        (ChainRulesCore.DoesNotExist(), ntuple((_->ChainRulesCore.DoesNotExist()), 0 + length(x))...)
                    end)
        end
    end
end

Macro output with this PR:

julia> Base.remove_linenums!(@macroexpand @scalar_rule sincos(x) cos(x) -sin(x))
quote
    if !(sincos isa ChainRulesCore.Type) && ChainRulesCore.fieldcount(ChainRulesCore.typeof(sincos)) > 0
        ChainRulesCore.throw(ChainRulesCore.ArgumentError("@scalar_rule cannot be used on closures/functors (such as $(sincos))"))
    end
    begin
        function (ChainRulesCore.ChainRulesCore).frule((ChainRulesCore._, var"##Δ1#277"), ::ChainRulesCore.typeof(sincos), x::Number)
            Ω = sincos(x)
            nothing
            return (Ω, ChainRulesCore.Composite{ChainRulesCore.typeof(Ω)}(cos(x) * var"##Δ1#277", -(sin(x)) * var"##Δ1#277"))
        end
    end
    begin
        function (ChainRulesCore.ChainRulesCore).rrule(::ChainRulesCore.typeof(sincos), x::Number)
            Ω = sincos(x)
            nothing
            return (Ω, begin
                        function sincos_pullback((var"##Δ1#278", var"##Δ2#279"))
                            $(Expr(:meta, :inline))
                            return (ChainRulesCore.NO_FIELDS, ChainRulesCore.muladd.(ChainRulesCore.conj(-(sin(x))), var"##Δ2#279", ChainRulesCore.:*.(ChainRulesCore.conj(cos(x)), var"##Δ1#278")))
                        end
                    end)
        end
    end
end

julia> Base.remove_linenums!(@macroexpand @non_differentiable length(x::Vector))
quote
    begin
        function (ChainRulesCore.ChainRulesCore).frule(::ChainRulesCore.Any, ::(Core).Typeof(length), x::Vector; var"##kwargs#280"...)
            return (length(x; var"##kwargs#280"...), ChainRulesCore.DoesNotExist())
        end
    end
    begin
        function (::(ChainRulesCore.Core).kwftype(ChainRulesCore.typeof(ChainRulesCore.rrule)))(var"##kwargs#281"::ChainRulesCore.Any, ::ChainRulesCore.typeof(ChainRulesCore.rrule
), ::(Core).Typeof(length), x::Vector)
            return (length(x; var"##kwargs#281"...), begin
                        function length_pullback(::ChainRulesCore.Any)
                            return (ChainRulesCore.DoesNotExist(), ChainRulesCore.DoesNotExist())
                        end
                    end)
        end
        function (ChainRulesCore.ChainRulesCore).rrule(::(Core).Typeof(length), x::Vector)
            return (length(x), begin
                        function length_pullback(::ChainRulesCore.Any)
                            return (ChainRulesCore.DoesNotExist(), ChainRulesCore.DoesNotExist())
                        end
                    end)
        end
    end
end

julia> Base.remove_linenums!(@macroexpand @non_differentiable length(x::Vector...))
quote
    begin
        function (ChainRulesCore.ChainRulesCore).frule(::ChainRulesCore.Any, ::(Core).Typeof(length), x::Vector...; var"##kwargs#282"...)
            return (length(x...; var"##kwargs#282"...), ChainRulesCore.DoesNotExist())
        end
    end
    begin
        function (::(ChainRulesCore.Core).kwftype(ChainRulesCore.typeof(ChainRulesCore.rrule)))(var"##kwargs#283"::ChainRulesCore.Any, ::ChainRulesCore.typeof(ChainRulesCore.rrule
), ::(Core).Typeof(length), x::Vector...)
            return (length(x...; var"##kwargs#283"...), begin
                        function length_pullback(::ChainRulesCore.Any)
                            return ChainRulesCore.ntuple(((::ChainRulesCore.Any,)->begin
                                            ChainRulesCore.DoesNotExist()
                                        end), 1 + ChainRulesCore.length(x))
                        end
                    end)
        end
        function (ChainRulesCore.ChainRulesCore).rrule(::(Core).Typeof(length), x::Vector...)
            return (length(x...), begin
                        function length_pullback(::ChainRulesCore.Any)
                            return ChainRulesCore.ntuple(((::ChainRulesCore.Any,)->begin
                                            ChainRulesCore.DoesNotExist()
                                        end), 1 + ChainRulesCore.length(x))
                        end
                    end)
        end
    end
end

This fixes #320.

Edit: Updated after the suggestions by the reviewers were incorporated.

codecov-io · 2021-03-31T11:12:52Z

Codecov Report

Merging #325 (1a64a05) into master (d675be3) will increase coverage by 0.06%.
The diff coverage is 94.11%.

@@            Coverage Diff             @@
##           master     #325      +/-   ##
==========================================
+ Coverage   89.76%   89.83%   +0.06%     
==========================================
  Files          13       13              
  Lines         469      472       +3     
==========================================
+ Hits          421      424       +3     
  Misses         48       48

Impacted Files	Coverage Δ
src/rule_definition_tools.jl	`96.12% <94.11%> (+0.09%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update d675be3...1a64a05. Read the comment docs.

mzgubic · 2021-03-31T11:20:44Z

Wow, that's great, thanks for looking into this! I will take a detailed look later today.

For now just a tip (I think from @oxinabox) I've found very useful when debugging macros: Base.remove_linenums!(@macroexpand @non_differentiable length(x::AbstractVector)) removes line numbers from the output and makes it much easier to read.

simeonschaub · 2021-03-31T11:45:43Z

src/rule_definition_tools.jl

+        function ($ChainRulesCore.rrule)($(esc_primal_sig_parts...))
            $(__source__)
-            return ($primal_invoke, $pullback_expr)
+            return ($(esc(primal_invoke)), $(pullback_expr))
        end


Since this is now escaped correctly, there shouldn't be a need to interpolate something like ChainRules anymore. That was just my proposed quick fix, since we didn't handle escaping correctly before.

simeonschaub · 2021-03-31T11:46:32Z

src/rule_definition_tools.jl

@@ -165,7 +165,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...)), ::typeof($f), $(inputs...))
+        function ($ChainRulesCore.frule)(($(esc(:_)), $(Δs...)), ::$(typeof)($f), $(inputs...))


No need to escape _

As mentioned in https://github.com/JuliaDiff/ChainRulesCore.jl/pull/325/files#r604881501, the main intention was to obtain a (somewhat) cleaner output. I don't mind changing it though.

Seems like _ getting mangled by macro hygiene is something that we should fix in Base. It definitely shouldn't get turned into a GlobalRef. Mind opening an issue?

Ok, JuliaLang/julia#40280 should fix this, so I think we don't need to escape it here.

oxinabox

Thanks, basically looks good too me.
Once you have addressed the comments to your satisfaction merge and tag a release

oxinabox · 2021-03-31T11:38:30Z

src/rule_definition_tools.jl

@@ -88,8 +88,8 @@ macro scalar_rule(call, maybe_setup, partials...)
    ############################################################################
    # Final return: building the expression to insert in the place of this macro
    code = quote
-        if !($f isa Type) && fieldcount(typeof($f)) > 0
-            throw(ArgumentError(
+        if !($f isa $Type) && $(fieldcount)($(typeof)($f)) > 0


How do you feel about not using brackets that are redudent?
On the one hand less visually noisy, on the other hand less consistent with other symbols interpolated in

Suggested change

if !($f isa $Type) && $(fieldcount)($(typeof)($f)) > 0

if !($f isa $Type) && $fieldcount($typeof($f)) > 0

I got rid of all the interpolations again: #325 (comment)

src/rule_definition_tools.jl

oxinabox · 2021-03-31T11:42:06Z

src/rule_definition_tools.jl

-        if !($f isa Type) && fieldcount(typeof($f)) > 0
-            throw(ArgumentError(
+        if !($f isa $Type) && $(fieldcount)($(typeof)($f)) > 0
+            $(throw)($(ArgumentError)(


TIL throw can be overwritten. It isn't a language keyword, just a built in function.
Though there seems like some inconsistency in what we are interpolating in.
In that we are interpolating in throw, but not ! or >, or isa.
All of which i think would be at least as common to shadow as throw (and still incredibly rare)

Same here, now macro hygiene takes care of all these things (the only downside being the messier output of @macroexpand).

oxinabox · 2021-03-31T11:54:27Z

test/rule_definition_tools.jl

-    # all that matters is that the following don't error, since they will resolve at
-    # parse time
-    using ChainRulesCore: ChainRulesCore
+using ChainRulesCore: @scalar_rule, @non_differentiable


BlueStyle doesn't have a rule on this, but I think that it should.
And that submodules within another file should be indented.
JuliaDiff/BlueStyle#85

Espeically in this case when there are so many of them nested, that i am having trouble following exactly what is going on.

I fixed the indentation.

oxinabox · 2021-03-31T11:56:06Z

test/rule_definition_tools.jl

+fixed(x) = :abc
+@non_differentiable fixed(x)
+
+# check name collision


Suggested change

# check name collision

# check name collision between a primal input called `kwargs` and the actual keyword args

I added your suggestion (unfortunately I missed it initially and couldn't commit the suggestion on Github anymore then).

oxinabox · 2021-03-31T11:57:24Z

src/rule_definition_tools.jl

        # Manually defined kw version to save compiler work. See explanation in rules.jl
-        function (::Core.kwftype(typeof(ChainRulesCore.rrule)))(kwargs::Any, rrule::typeof(ChainRulesCore.rrule), $(primal_sig_parts...))
-            return ($(_with_kwargs_expr(primal_invoke)), $pullback_expr)
+        function (::$(Core.kwftype)($(typeof)($(rrule))))($(esc(kwargs))::$(Any), ::$(typeof)($(rrule)), $(esc_primal_sig_parts...))


I am not sure you can actually shadow Core

Again this should be handled by macro hygiene.

oxinabox · 2021-03-31T11:58:44Z

test/rule_definition_tools.jl

+my_id(x) = x
+@scalar_rule(my_id(x), 1.0)
+
+module IsolatedSubmodule


what is being tested here?

That all frule and rrules can be called without errors - this would detect e.g. if the macro output would try to look up something in ChainRulesCore in the module the rules were defined in (i.e., in IsolatedModuleForTestingScoping.ChainRulesCore). To avoid having to import anything from ChainRulesCore in addition to @scalar_rule and @non_differentiable, these checks are performed in a separate submodule.

Can we indicate this in the code with a comment?

mzgubic · 2021-03-31T12:50:59Z

src/rule_definition_tools.jl

-    )
-    return esc(@strip_linenos quote
+    pullback_expr = @strip_linenos quote
+        function $(esc(propagator_name(primal_name, :pullback)))($(esc(:_)))


Thanks for changing this as well, it's much clearer now. ~~Similarly here I don't think we need to escape :_~~

Yeah, escaping :_ ends up with (arguably) cleaner output but makes the implementation a bit annoying to read.

simeonschaub

Have you tried this without all the interpolation, just with the changes to how esc is used? I would prefer being careful with how we use esc instead of just interpolating everything, since that's really how macro hygiene is intended to be used and I find the interpolations make the source code harder to read. I also wouldn't be too worried about the output of @macroexpand, as long as stacktraces are good.

devmotion · 2021-03-31T13:35:31Z

Have you tried this without all the interpolation, just with the changes to how esc is used? I would prefer being careful with how we use esc instead of just interpolating everything, since that's really how macro hygiene is intended to be used and I find the interpolations make the source code harder to read. I also wouldn't be too worried about the output of @macroexpand, as long as stacktraces are good.

I agree, I already started to reduce the numbers interpolations in my local branch since I had the same feeling. I was quite happy about the (arguably) clean output of @macroexpand but I don't think it's worth it anymore.

simeonschaub · 2021-03-31T15:10:54Z

src/rule_definition_tools.jl

@@ -356,7 +357,7 @@ function tuple_expression(primal_sig_parts)
    else
        num_primal_inputs = length(primal_sig_parts) - 1 # - vararg
        length_expr = :($num_primal_inputs + length($(esc(_unconstrain(primal_sig_parts[end])))))
-        Expr(:call, :ntuple, Expr(:(->), :_, DoesNotExist()), length_expr)
+        Expr(:call, :ntuple, Expr(:(->), :($(esc(:_))), DoesNotExist()), length_expr)


We should not rely on this as a fix, since this behavior will likely change soon.

What do you suggest? Using ::Any instead? Or just :_? I also noticed that if esc is removed both in the pullback function signature and this anonymous function, then the macro hygiene will replace both occurrences of _ with the same variable name (something like `var"#249#_").

Ah, in that case I'd go with ::Any or just something like _tmp.

simeonschaub

Just this small nit plus removing all remaining uses of esc(:_) and this looks good to go. Sorry for being so picky here, but I am now very happy with this. Macro hygiene is always a bit difficult to get right, since mistakes can show up in subtle ways or only in edgecases. Thanks for sticking with me!

simeonschaub · 2021-03-31T17:04:58Z

src/rule_definition_tools.jl

        # Manually defined kw version to save compiler work. See explanation in rules.jl
-        function (::Core.kwftype(typeof(ChainRulesCore.rrule)))(kwargs::Any, rrule::typeof(ChainRulesCore.rrule), $(primal_sig_parts...))
-            return ($(_with_kwargs_expr(primal_invoke)), $pullback_expr)
+        function (::Core.kwftype(typeof($rrule)))($(esc(kwargs))::Any, ::typeof($rrule), $(esc_primal_sig_parts...))


Interpolation of rrule can be removed?

simeonschaub

Awesome, thanks!

devmotion · 2021-03-31T17:36:39Z

Would it be useful to add @nospecialize to the discarded argument of type Any in the frule and in the pullback of rrule generated by @nondifferentiable (I guess it doesn't matter for the ntuple expression and can't be applied to the tuple destructuring in the frule of @scalar_rule)?

simeonschaub · 2021-03-31T17:40:59Z

Yes, I think that would make sense, but probably better as a separate PR.

mzgubic · 2021-03-31T22:31:30Z

all good from my side, @simeonschaub and @oxinabox are far more knowledgeable than me anyway, I wanted to review to learn really

devmotion added 2 commits March 31, 2021 12:14

Interpolate everything in rule helpers

6dac9a7

Bump version

1a64a05

simeonschaub reviewed Mar 31, 2021

View reviewed changes

oxinabox approved these changes Mar 31, 2021

View reviewed changes

mzgubic reviewed Mar 31, 2021

View reviewed changes

simeonschaub requested changes Mar 31, 2021

View reviewed changes

devmotion added 3 commits March 31, 2021 16:34

Less interpolation, simpler code, messier @macroexpand output...

4571cf4

Change indentation of submodules in tests

cf6d283

Work around JuliaLang/julia/issues/32727

c687250

simeonschaub reviewed Mar 31, 2021

View reviewed changes

Add additional comment

f4fcdf2

devmotion changed the title ~~Interpolate everything in rule helpers~~ Interpolate ~~everything~~ more in rule helpers and fix escaping of @non_differentiable Mar 31, 2021

devmotion changed the title ~~Interpolate ~~everything~~ more in rule helpers and fix escaping of @non_differentiable~~ Interpolate more in rule helpers and fix escaping of @non_differentiable Mar 31, 2021

simeonschaub reviewed Mar 31, 2021

View reviewed changes

devmotion added 3 commits March 31, 2021 19:20

Remove remaining esc(:_)

3f474ee

Remove unneeded interpolation of rrule

28ae5e6

Add explanations to tests

c083278

simeonschaub approved these changes Mar 31, 2021

View reviewed changes

devmotion requested a review from mzgubic March 31, 2021 17:58

devmotion merged commit fb12855 into JuliaDiff:master Mar 31, 2021

devmotion deleted the dw/macro_hygiene branch April 1, 2021 07:03

devmotion mentioned this pull request Apr 1, 2021

Mark discarded arguments with @nospecialize #326

Merged

	if !($f isa $Type) && $(fieldcount)($(typeof)($f)) > 0
	if !($f isa $Type) && $fieldcount($typeof($f)) > 0

	# check name collision
	# check name collision between a primal input called `kwargs` and the actual keyword args

Interpolate more in rule helpers and fix escaping of @non_differentiable #325

Interpolate more in rule helpers and fix escaping of @non_differentiable #325

Conversation

devmotion commented Mar 31, 2021 • edited Loading

codecov-io commented Mar 31, 2021 • edited Loading

Codecov Report

mzgubic commented Mar 31, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

simeonschaub Mar 31, 2021 • edited Loading

Choose a reason for hiding this comment

oxinabox left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mzgubic Mar 31, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

simeonschaub left a comment

Choose a reason for hiding this comment

devmotion commented Mar 31, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

simeonschaub left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

simeonschaub left a comment

Choose a reason for hiding this comment

devmotion commented Mar 31, 2021

simeonschaub commented Mar 31, 2021

mzgubic commented Mar 31, 2021

Interpolate more in rule helpers and fix escaping of `@non_differentiable` #325

Interpolate more in rule helpers and fix escaping of `@non_differentiable` #325

devmotion commented Mar 31, 2021 •

edited

Loading

codecov-io commented Mar 31, 2021 •

edited

Loading

simeonschaub Mar 31, 2021 •

edited

Loading

mzgubic Mar 31, 2021 •

edited

Loading

simeonschaub left a comment •

edited

Loading