For a particular JAX function func, one can define non-differentiable arguments by using the decorator @partial(jax.custom_jvp, nondiff_argnums=...). However, in order to make it work, one must also explicitly define the differentiation rules in a custom jvp function by using the decorator @func.defjvp. I'm wondering if there is a generic way to define non-differentiable argument for any given func, without defining a custom jvp (or vjp) function? This will be useful when the differentiation rules are too complicated to write out. Continue reading...
