### Author:  [@oxinabox](https://discourse.julialang.org/u/oxinabox)
#### Post date: [May 15, 2024, 5:06am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/2 "2024-05-15T05:06:04Z")
So ChainRules doesn’t support giving deriviatives wrt to keyword arguments.
Only receiving the keyword argument primals.
Since it was found basically they are never needed, because they tend to be things like
or they tend to dispatch to something with positional arguments only under the hood.
The normal work around is to rewrap the thing into something that takes the keyword arguments in question as positional and write a rule for that.
However, it _is_ possible.
Wether or not a particular AD supports it I can’t promiste.
For some function `foo(x; a, b)`
The following works for frules for Differactor:
```julia
function ChainRulesCore.frule((_, dkwargs, _, dx), ::typeof(kwcall), kwargs, ::typeof(foo), x)
y = foo(x; kwargs...)
dy = ...
return y, dy
end
```
and is used eg in [CedarEDA.jl/src/extra\_rules.jl at 759d5c5c173f509829fb741e748ecffd85303ec5 · CedarEDA/CedarEDA.jl · GitHub](https://github.com/CedarEDA/CedarEDA.jl/blob/759d5c5c173f509829fb741e748ecffd85303ec5/src/extra_rules.jl#L59-L61)
For reverse mode the equivalent is
```julia
function ChainRulesCore.rrule(::typeof(kwcall), kwargs, ::typeof(foo), x)
y = foo(x; kwargs...)
function foo_pullback(dy)
da = haskey(kwargs, :a) ? ... : NoTangent()
db = haskey(kwargs, :b) ? ... : NoTangent()
dkwargs = Tangent{typeof(kwargs)}(a=da, b=db)
dx = ...
return NoTangent(), dkwargs, NoTangent(), dx
end
end
```
In particular I am not sure Zygote will work right with this, it might, it might not.
I know this line
>
drops the derivatives from keyword arguments but I am not sure if it will actually be hit if you provide a overload for the `rrule` for `kwcall` directly.
---
### Author:  [@RainerHeintzmann](https://discourse.julialang.org/u/RainerHeintzmann)
#### Post date: [May 15, 2024, 6:53am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/3 "2024-05-15T06:53:08Z")
Thanks a lot for your reply! I think `kwargs` support would be really nice and I am sure that there are many use cases. E.g. in `IndexFunArrays.jl` or `SeparableFunctions.jl` (not realeased) the offset, center and scale of each function can be modified from the default via keyword arguments. In my current usecase I want to have AD support for `SeparableFunctions.jl` to get a fast Gaussian fitting routine where typically one would use the `pos` keyword argument.
Yes, trying to define the rule for the `kwcall` function is what I tried as well, but so far without any success:
```julia
using ChainRulesCore
using Zygote
foo(a,b;c)=a*b*c
foo(1,2;c=3) # works
g = Core.kwcall((c = 3,), foo, 1, 2) # works
function ChainRulesCore.rrule(::typeof(Core.kwcall), kwargs, ::typeof(foo), a, b)
println("kwargs foo in chainrule")
y = foo(a,b; kwargs...)
function foo_pullback(dy)
println("kwargs foo pulling back")
da = dy*b*c
db = dy*a*c
dc = haskey(kwargs, :c) ? dy*a*b : NoTangent()
dkwargs = Tangent{typeof(kwargs)}(c=dc)
return NoTangent(), dkwargs, NoTangent(), da, db
end
end
gradient((x)->foo(2.0,3.0,c=x), 1.0) # works but does not use the above chain rule
```
as one can see the chain rule is never used by the `gradient` function, but it is nevertheless able to deal with `kwargs`. The interesting definition you found in `Zygote.jl` seems to ignore kwargs, but the framework somehow does suppports it.
Since there does not seem to be a way to specify kwargs support for gradients in `CainRulesCore.jl`, maybe it would be time to come up with one?
E.g. One could wrap the return type in `KWArgsTangent((standardTangents...),(NamedTangents...))` or something like this?
But would this be type stable?
---
### Author:  [@oxinabox](https://discourse.julialang.org/u/oxinabox)
#### Post date: [May 15, 2024, 9:28am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/4 "2024-05-15T09:28:43Z")
I would be most tempted just to change Zygote so the code above worked.
Its nonambigious, though not the nicest.
And I think probably it could be made to work.
It would just need to be changed to check if there was an actual `rrule` for `kwfunc` before falling back to just calling the rule for the function with primal keywords.
And something like it would be required in order to support anything nicer anyway.
Maybe for the something nicer we can gate the rule on a conflig trait type parameter, like we do for the calling back into AD traits.
but first would definately need to make the above work in Zygote cos that is where the gradients need to end up.
As a workaround, I think for Zygote only defining
```julia
using ZygoteRules
function ZygoteRules._pullback(::typeof(Core.kwcall), kwargs, ::typeof(foo), a, b))
println("kwargs foo in zygote rule")
y = foo(a,b; kwargs...)
function foo_pullback(dy)
println("kwargs foo pulling back")
da = dy*b*c
db = dy*a*c
dc = haskey(kwargs, :c) ? dy*a*b : nothing()
dkwargs = (;c=dc)
return nothing, dkwargs, nothing da, db
end
end
```
should work.
It should get hit before the chainrule’s inferface gets hit
---
### Author:  [@RainerHeintzmann](https://discourse.julialang.org/u/RainerHeintzmann)
#### Post date: [May 15, 2024, 10:01am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/5 "2024-05-15T10:01:34Z")
You are right. Changing Zygote to be able to deal with `kwcall` `rrule` definitions should be the first step. Yet Zygote does not even seem to deal with your suggestion (a few corrections):
```julia
using ZygoteRules
foo(a,b;c)=a*b*c
function ZygoteRules._pullback(::typeof(Core.kwcall), kwargs, ::typeof(foo), a, b)
println("kwargs foo in zygote rule")
y = foo(a,b; kwargs...)
function foo_pullback(dy)
println("kwargs foo pulling back")
da = dy*b*c
db = dy*a*c
dc = haskey(kwargs, :c) ? dy*a*b : NoTangent()
dkwargs = (;c=dc)
return nothing, dkwargs, nothing, da, db
end
end
using Zygote
gradient((x)->foo(2.0,3.0,c=x), 1.0) # does not call the above definition
```
Does the `rrule` for `kwcall` work directly with `Diffractor.jl`?
---
### Author:  [@oxinabox](https://discourse.julialang.org/u/oxinabox)
#### Post date: [May 15, 2024, 10:08am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/6 "2024-05-15T10:08:43Z")
Right now rrules with keyword arguments don’t work in Diffractor at all.
Though there is a PR that fixes that and makes the above rrule for kwcall work.
For Zygote i am pretty sure that should work.
I would need to dig in and debug it to workout why it doesn’t.
---
### Author:  [@oxinabox](https://discourse.julialang.org/u/oxinabox)
#### Post date: [May 20, 2024, 6:10am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/7 "2024-05-20T06:10:35Z")
Oh I relalised I had the signature of `_pullback` wrong, forgot the context argument.
and there were a few other issues.
I have fixed them and the following works. 🎉
```julia
using ZygoteRules
foo(a,b;c=42.0)=a*b*c
function ZygoteRules._pullback(::ZygoteRules.AContext, A::typeof(Core.kwcall), kwargs, ::typeof(foo), a, b)
println("kwargs foo in zygote rule")
y = foo(a,b; kwargs...)
c = get(kwargs, :c, 42.0) # insert default as we need that in pullback.
function foo_pullback(dy)
println("kwargs foo pulling back")
da = dy*b*c
db = dy*a*c
dc = haskey(kwargs, :c) ? dy*a*b : NoTangent()
dkwargs = (;c=dc)
return nothing, dkwargs, nothing, da, db
end
return y, foo_pullback
end
using Zygote
gradient((x)->foo(2.0,3.0,c=x), 1.0)
```
I made that a little more complex so you could see how to handle kwarg with a default.
---
### Author:  [@RainerHeintzmann](https://discourse.julialang.org/u/RainerHeintzmann)
#### Post date: [May 20, 2024, 10:09am UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/8 "2024-05-20T10:09:26Z")
Thanks! This is a nice step forward. But then there is still the problem that this solution would be Zygote-specific and there seem to be no way to specify this as an `rrule`. I also assume that ZygoteRules is a heavier package than ChainRulesCore, and in each case there is a `kwargs` involved for which AD is wanted, one would need to make the package depend on it.
---
### Author:  [@oxinabox](https://discourse.julialang.org/u/oxinabox)
#### Post date: [May 20, 2024, 3:57pm UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/9 "2024-05-20T15:57:32Z")
No ZygoteRules is no heavier than ChainRulesCore.
In fact, I expect it is technically ligher, but CRC is already so fast to load it doesn’t matter.
Its just that it is Zygote specific, and doesn’t work for any of the other 9 ADs that support ChainRulesCore (to varying extents).
There is indeed no way to specify this as a rrule that Zygote will recognize.
As I said, I think what we need to do is teach Zygote to recognize this in `rrule` form.
And _then_ maybe add some syntactic sugar to do this in ChainRulesCore without the `kwfunc` weirdness. Maybe.
The code in Zygote that would need to be changed is:
>
and
>
Which is not the most complicated code in the world.
But it isn’t simple.
---
### Author:  [@RainerHeintzmann](https://discourse.julialang.org/u/RainerHeintzmann)
#### Post date: [July 3, 2024, 3:05pm UTC](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150/10 "2024-07-03T15:05:34Z")
There seem to be other problems. I tried implementing your idea for my problem, but failed when the datastructures are a little more complex:
```julia
using Zygote, ZygoteRules
foo(a; c=[42.0, 43.0]) = (a .*c, )
function ZygoteRules._pullback(::ZygoteRules.AContext, A::typeof(Core.kwcall), kwargs, ::typeof(foo), a)
println("kwargs foo in zygote rule")
y = foo(a; kwargs...)
c = get(kwargs, :c, [42.0, 43.0]) # insert default as we need that in pullback.
function foo_pullback(dy)
println("kwargs foo pulling back: dy=$dy, c=$c")
da = (c' * dy[1],) # (reshape(dy .*c[1], (2,1)), reshape(dy .*c[2], (1,2))) # dy .*c
dc = haskey(kwargs, :c) ? (a * dy[1],) : NoTangent() # # (reshape(dy[1] .*a, (2,1)), reshape(dy[2] .*a, (1,2)))
dkwargs = (;c=dc)
@show da
return nothing, dkwargs, nothing, da
end
return y, foo_pullback
end
q = pullback((x)->foo(x), 1.0)[2](([1.0, 1.0],)) # works fine and returns (85.0,)
gradient((x)->sum(.*(foo(x)...)), 2.0) # works fine!
q = pullback((x)->foo(x, c=[42.0, 43.0]), 1.0)[2](([1.0, 1.0],)) # yields an error. WHY? How to fix?
gradient((x)->sum(.*(foo(x, c=[42.0, 43.0])...)), 2.0) # yields an error. WHY?
```
The error is
`ERROR: MethodError: no method matching (::ChainRulesCore.ProjectTo{Float64, @NamedTuple{}})(::Vector{Float64})`
Since the returned result is identical in both cases and one works and the other does not, I am a bit puzzled. It seems like Zygote is doing something else behind the scenes which may need fixing to make this idea work?