Rules for functions / structs etc

[willtebbutt]

The current design requires that rules have the form

function rrule(::typeof(foo), args...)
    ...
    return (rule_arg_1, rule_arg_2, ...)
end

I would like to propose to extend this to the following:

function rrule(foo::typeof(foo), args...)
    ...
    return (rule_foo, rule_arg_1, rule_arg_2, ...)
end

A lot of the time rule_foo will be a DNE, which is fine. However, if we're looking to implement rules for particular structs, then it may well be important to be able to represent the gradient w.r.t. the fields of the struct.

[ChrisRackauckas]

Won't this be necessary for things like the derivative of Dense?

[oxinabox]

I am pretty sure this is actually closed.
Pretty sure we determined
The derivative of something returing a struct is a named tuple of differentiables, 1 per field.

@willtebbutt ?

[oxinabox]

Won't this be necessary for things like the derivative of Dense?

Not under normal uses no.
Tracker can't work on things returning structs (I do believe Nabla and Zygote can).
the derivative of a Flux model is not for Dense, but rather it is params(::Dense) which is a collection of matrixs.

[willtebbutt]

Won't this be necessary for things like the derivative of Dense?

There are two orthogonal issues here:

1. If you're using an AD system that can handle adjoints w.r.t. structs (e.g. Zygote) and you wanted to define a custom rrule for Dense as opposed to just running AD then yes, you definitely need this.
2. Not all AD systems are able to handle rules of this form - as @oxinabox points out Tracker doesn't, and Nabla only kind of does. For these systems you have to write work-arounds that pull the data out of structs first and then pass to a plain function, so it's not so much of a problem.

My feeling is that Zygote has got this 'right', and that we should support this functionality. Unfortunately this change requires some fairly dull work.

[oxinabox]

The only remaining parts of this are covered by #8