@non_differentiable

[willtebbutt]

Similarly to @nograd in Zygote, we need a @non_differentiable macro to make it straightforward for someone to declare that the particular method of a function should have an frule and rrule that return zeros / related.

Proposed Usage:

@non_differentiable f(x::Tx, y...)

should be expanded to something like

function frule((::Any, ::Any, ::Vararg), ::typeof(f), x::Tx, y...)
    return f(x, y...), Zero()
end

function rrule(::typeof(f), x::Tx, y...)
    function pullback_f_non_differentiable(::Any)
        return (DoesNotExist(), DoesNotExist(), map(DoesNotExist(), 1:length(y)))
    end
    return f(x, y...), pullback_f_non_differentiable
end

I think the frule and rrule that something should expand to once we've established the API is pretty clear (I'm not sure I've actually got the rrule correct for the vararg in the above). Same thing for what the kwargs bit of the API should look like.

So the point of discussion is does the above form of @non_differentiable do what we want it to do? In particular, do we want to insist / not insist that people provide argument names?

[cossio]

Why not a shorter name, like @nograd?

[willtebbutt]

My feeling is that @non_differentiable is more descriptive of what's going on.

To my ear @nograd suggests feels a bit specific to single-output functions, and it's not really clear what it means ie. why is there no gradient?

[cossio]

True that grad is specific to single-output functions.
Maybe @nodiff?

It would be nice to also have a macro that shields a block of code from being differentiated (instead of just a function).

[willtebbutt]

It would be nice to also have a macro that shields a block of code from being differentiated (instead of just a function).

Hmm yeah, this would be nice. Might be possible to do at the CR level with a closure... not sure if we would run into world-age issues or something though (as in, I really have no idea whether this is doable at the CR level or whether it has to be done at the AD level)

[simeonschaub]

Shouldn't that be the AD tool's job though, not ChainRules'? I don't really see, how something like that would go into our rule definitions.

[willtebbutt]

You could just see it as a special case of a user-defined rule, and we generally want users to be able to define custom rule and expect AD to play nicely with them, regardless which AD tool they use.

edit: reworded

[oxinabox]

I wonder if we can automatically determine which should be DoesNotExist and which should be Zero.

e,g, we know that Symbols, AbstractStrings and Char have corresponding differentiable of DoesNotExist.
for everything else Zero i guess.
Alt could introduct a NonSpecificNonexistance <: AbstractZero

[MikeInnes]

Probably worth having a look at Zygote's utils file. Things like dropgrad, ignore (@cossio's suggestion), hook and @showgrad could all be easily incorporated into CR and at least ignore and hook seem sufficiently useful to do so. [Some utils have the name grad in them for historical reasons but work just as well in forward and reverse in principle.]

They are really easy to define (using Zygote syntax here):

hook(f, x) = x
@tangent hook(f, x) = x, ẋ -> f(ẋ)
@adjoint hook(f, x) = x, x̄ -> (nothing, f(x̄))

ignore(block) = block()
@tangent hook(block) = block(), _ -> nothing
@adjoint hook(block) = block(), _ -> nothing

ignore() do # can also use a thin macro layer like `@ignore begin ...`
  @info "ignored logging code"
end

hook(dx -> @show(x, dx), x) # print x along with its gradient/derivative

Zygote and co could then re-export these.

FYI, @nograd as it exists in Zygote is kind of a hack from very early in Zygote's existence and should probably be removed at some point. It's way to easy too define over-general rules and so better to encourage people to use signatures in general. In the @nograd case we had issues with things like rand which are non-differentiable in their Base form, but packages add differentiable methods. This proposal is already well-aligned with that but I just wanted to make it explicit.

[willtebbutt]

While I remember: this would also be a great opportunity to implement a number of very specific rules for higher-order functions. For example, if foo isn't differentiable, you know that map(foo, data) also isn't.

[Keno]

👍 On this. Seeing lots of unnecessary work being done by my new AD pass because it's trying to work on things that definitely don't need to be AD'd.