Closed shoyer closed 1 month ago
the same trick of using "custom ops" to wrap entire diffable programs works for e.g. wrapping TF into a PyTorch op. I agree it would be nice to add a clean API for that.
Here's a working version of a prototype for passing gradient between TensorFlow 2 and JAX, which may be a useful point of reference: https://gist.github.com/shoyer/5f72853c2788e99e785f4737ee8a6ae1
Thanks for the suggestion, but I think we'll leave this to user-defined wrappers rather than adding it to the JAX public API.
There's a nice duality between
custom_gradient
andjax.jvp
/tf.gradients
that is exactly what is needed for piping gradients back and forth, e.g.,It would be nice to expose user facing functions for this.
Note that defining custom gradients rules isn't quite enough on its own. You also need to define a new JAX "primitive" in order to wrap non-traceable computations.