Jax Custom Ops & Calling TF Custom Ops Using jax2tf.call_tf

[adam-hartshorne]

I have a number of TF custom ops that call c++ code (which themselves depend on some c++ libraries) and have their own associated custom gradients, which I have built as outlined here

https://www.tensorflow.org/guide/create_op

In theory does this new functionality of jax2tf.call_tf support calling such TF ops?

And is there any immediate plans to create similar functionality to the TF custom op functionality for JAX? I would really love to be able to be able to interface with c++ libraries from within JAX framework.

[jekbradbury]

I don't think custom TF ops will work through call_tf (which is currently implemented by going through XLA). But you can call them from Python code during jitted function execution using host_callback or directly from jitted JAX code by registering a custom XLA op (a great third-party tutorial for that is here: https://github.com/dfm/extending-jax).