A concise way to define a custom jvp rule for a wrappers with many inputs and defaults?

[lankef]

Hello folks!

I have a jitted function with about 40 inputs, fg(), with less than ten differentiable inputs. It implements the value and gradient of a function f. I want to split it into two wrappers, f() and g() that take the same inputs. f will return the value only, and g will be the custom_jvp of f. I am currently copying and pasting all of the inputs and defaults from fg to f, but this is becoming very lengthy. It can also break easily if someone decided to add new parameters to fg. Is there a good way to handle this?

To my understanding the @custom_jvp flag doesn't work with **kwargs, so I can't just define f using fg(*args, **kwargs). This will cause TypeError: keyword arguments (['stuff']) could not be resolved to positions...

[dfm]

Good question! I'm not sure I totally follow the motivation for the change you're trying to make. Why do you need these two frontend functions? Could you post some explicit example code to guide the discussion?

At a high level, one option would be to have only one function accept the explicit keywords and then pack all the non-differentiable arguments into a config or params dict that is then passed to the lower level functions.

[lankef]

It's a nested iteration that I have already implemented implicit derivatives for. The problem with using dict is that dicts are non-hashable, and cannot contain static arguments.

[dfm]

I don't think that the non-differentiable arguments to custom_jvp need to be hashable! But, regardless, you could use a different data structure (NamedTuple, frozen dataclass, etc.). It's hard to provide any specific advice without a clearer worked example!

[lankef]

I see. You're right. I have another question... What is the best way, then, to stop jax from differentiating config?

@custom_jvp
def f(x, y, config):
    print(config['test'])
    return jnp.sin(x) * y + config['constant']

@f.defjvp
def f_jvp(primals, tangents):
    x, y, config = primals
    x_dot, y_dot, _ = tangents
    primal_out = f(x, y, config)
    tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
    return primal_out, tangent_out

Executing grad(f)(1., 2., {'test': 'good', 'constant':0.1}) or grad(partial(f, config={'test': 'good', 'constant':0.1}))(1., 2., ) will throw a TypeError: Value 'good' with type <class 'str'> is not a valid JAX type. Since config is a dict and non-hashable, we can't set it to static. We also can't use stop_gradient.

[dfm]

In this case, you want to use the nondiff_argnums argument to custom_jvp:

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.custom_jvp, nondiff_argnums=(2,))
def f(x, y, config):
  print(config['test'])
  return jnp.sin(x) * y + config['constant']

@f.defjvp
def f_jvp(config, primals, tangents):
  x, y = primals
  x_dot, y_dot = tangents
  primal_out = f(x, y, config)
  tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
  return primal_out, tangent_out

jax.grad(f)(1., 2., {'test': 'good', 'constant':0.1})

Note that this also changes the signature of f_jvp. Does that do the trick in your case?

[lankef]

Yup! Thanks alot!