Keyword arguments not working with @jax.custom_gradient

[AlexeyG]

@jax.custom_gradient
def my_fun(x, y, c=1.):
  return c * (x + y), lambda x, y: (c * (1. + y), c * (x + 1.))

f = lambda x, y: jnp.square(my_fun(x, y, c=2.)).sum()

Calling f(10., 5.) gives

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-55-de80b4d2e2a5> in <module>()
----> 1 f(10., 5.)

<ipython-input-54-d90a45f76782> in <lambda>(x, y)
      3   return c * (x + y), lambda x, y: (1. * c, 1. * c)
      4 
----> 5 f = lambda x, y: jnp.square(my_fun(x, y, c=2.)).sum()

TypeError: __call__() got an unexpected keyword argument 'c'

[mattjj]

I'm working on a custom_gradient overhaul, and I made sure that my prototype passes this test. I hope to have a PR (i.e. a branch you can work off of) this week, and maybe something mergeable next week.

[AlexeyG]

Sounds great. Thanks a lot, @mattjj !

[mattjj]

#2026 changed the API for custom derivatives, and also added support for default argument values and keyword arguments!