# Problems with AD inference on wrapper function

**URL:** https://discourse.julialang.org/t/problems-with-ad-inference-on-wrapper-function/116454
**Category:** General Usage
**Tags:** question
**Created:** [July 1, 2024, 4:58am UTC](https://discourse.julialang.org/t/problems-with-ad-inference-on-wrapper-function/116454 "2024-07-01T04:58:04Z")
**Posts on this page:** 1
**Showing post:** 4

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [July 1, 2024, 11:28pm UTC](https://discourse.julialang.org/t/problems-with-ad-inference-on-wrapper-function/116454/4 "2024-07-01T23:28:01Z")

</div>

The answer is relatively straightforward but does require some background on how keyword function dispatch works under the hood, [Julia Functions · The Julia Language](https://docs.julialang.org/en/v1.10/devdocs/functions/#Keyword-arguments). As seen from the example, the auto-generated “keyword sorter” function (a method of `Core.kwcall` on newer versions of Julia) uses conditionals to handle the presence or absence of certain kwargs. Because Zygote will unconditionally generate type unstable code when it encounters branching control flow, this means all calls with keyword arguments will be type unstable under AD.

As you’ve noted, this can be worked around by defining an rrule for the function in question. This is how e.g. `sum(...; dims=...)` can be type stable with Zygote: there are rules for it in ChainRules.jl. Unfortunately, having an rrule does preclude differentiating wrt. keyword arguments. @oxinabox has a far more in-depth series of posts about this and possible workarounds at [Rrule (or frule) with kwargs](https://discourse.julialang.org/t/rrule-or-frule-with-kwargs/114150).

---

_[View the full topic](https://discourse.julialang.org/t/problems-with-ad-inference-on-wrapper-function/116454)._