Composing with an inactive function

[cgeoga]

Apologies if this has been discussed in an issue and I've just missed it in my search.

Say you have a function h(x) = g(f(x)), and for some reason the best way to compute h'(x) is with g(f'(x)). I have the following working example of that workflow (obviously the derivatives themselves are incorrect, so please ignore that part) that uses the very enjoyable custom rules:

using Enzyme
using Enzyme.EnzymeRules

basic_inner(x) = cos(x)
basic_outer(x) = sin(x)

# Add what I would have naively called an "inactivity" rule for basic_outer.
function EnzymeRules.forward(func::Const{typeof(basic_outer)},
                             ::Type{<:Duplicated},
                             x::Duplicated)
  Duplicated(basic_outer(x.val), x.dval)
end

# a simple composed function.
basic(x) = basic_outer(basic_inner(x))

# I want this to give back (cos(1.0), -sin(1.0)), not (cos(1.0), 0.0), which is
# what I get if I define EnzymeRules.inactive(::typeof(basic_outer), args...)=nothing.
println(autodiff(Forward, basic, Duplicated, Duplicated(1.0, 1.0)))

This is great, and it gives me what I want. But in more complex applications where functions like basic_outer take more arguments, or sometimes you want different kinds of derivatives and so on, I could imagine this being slightly more annoying. I think I might just not understand the concept of EnzymeRules.inactive, but it would be super nice if there were some way to just define a method like

EnzymeRules.passive(::typeof(basic_outer), args...) = nothing # or something, I just made up the word passive here

that sort of just propagates the dval like my manual method does here. Is something like that conceivably possible? Does it already exist and I've just missed it in some docs?

Thanks so much for reading and for working on such an exciting library.

[wsmoses]

To confirm, you want to override the differentiation of the method to behave like the identity function?

[cgeoga]

That's right, yeah. Sorry, I see that that wasn't super clear on reading again. But that's exactly what I'd like. The strategy with the specific forward rules does work for this, I'm just curious if there is some more elegant way to create that behavior.

[wsmoses]

I'm not sure I see why this is likely to be common. Mind explaining the uses?

[cgeoga]

That's fair. I encountered the need for something like this playing with differentiating Bessel functions with respect to order at some edge cases, where in our case the f(x) computed a bunch of partial sum terms and g(x) was a series acceleration. I wouldn't be entirely surprised if this came up in other aspects of numerics at that low level, but I don't have a great list.

With that said though, I know you're triaging your time a bit and we actually got to a solution using extensions plus the custom rules, which I really enjoy the design of after using them a bit (this PR, if you are curious). In the end, the best solution doesn't even exactly use what I ask for here. So I think it might be best to close this and maybe I'll reopen it if I ever end up with a more compelling use-case or a longer list of examples where it would be useful. Feel free to close or let me know that you're cool with me doing so.

Good luck with the defense, by the way! I've seen you mention it in a discourse thread or something.