feature request: Specifying non-differentiable outputs in custom_vjp

Currently the custom_vjp interface allows one to specify input arguments that should not be differentiated via the nondiff_argnums argument.

It would be useful to have an equivalent option for a function with multiple return values, allowing the user to specify which output arguments should not be differentiable. Perhaps it could be named nondiff_outputs?

[proteneer]

Sorry for the late reply, and thank you for this. I’ve verified that it works as intended.

[ianwilliamson]

Something along these lines:

import dataclasses
import jax
import jax.numpy as jnp

@dataclasses.dataclass
class Foo:
   jax_thing: jnp.array
   non_jax_thing: int
   
def flatten(foo):
  return (foo.jax_thing,), (foo.non_jax_thing,)

def unflatten(aux_data, children):
  non_jax_thing, = aux_data
  jax_thing, = children
  return Foo(jax_thing, non_jax_thing)

jax.tree_util.register_pytree_node(Foo, flatten, unflatten)

If you had nothing differentiable in Foo you can just return an empty tuple as the first output of flatten() and ignore the second arg of unflatten().