jax cpp interface

[mgbukov]

Is there (any work towards) a c++ interface for jax?

[skye]

What would you like to do exactly? It's currently possible to save a jit'd version of a function as an XLA program, which can then be invoked from C++ via XLA's C++ client. Check out jax/tools/jax_to_hlo.py (HLO is how XLA programs are represented).

[mgbukov]

I was thinking of building and evaluating networks on a GPU and computing gradients quickly.

…

On Mon, Dec 16, 2019 at 1:37 PM Skye Wanderman-Milne < ***@***.***> wrote: What would you like to do exactly? It's currently possible to save a jit'd version of a function as an XLA program, which can then be invoked from C++ via XLA's C++ client <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/client/local_client.h>. Check out jax/tools/jax_to_hlo.py <https://github.com/google/jax/blob/master/jax/tools/jax_to_hlo.py> (HLO is how XLA programs are represented). — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#1871?email_source=notifications&email_token=ADHSUBTHQY735HAGU7A6PNTQY7YJNA5CNFSM4J3PNNZKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEHAF2RQ#issuecomment-566254918>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADHSUBWTPUWYFO5NFLJAMR3QY7YJNANCNFSM4J3PNNZA> .

-- Marin Bukov, PhD Gordon and Betty Moore postdoctoral fellow, Dept. of Physics University of California, Berkeley 366 LeConte Hall MC 7300, Berkeley, CA 94720-7300 url: https://mgbukov.github.io/

[skye]

With the above tooling I mentioned, you could build the network + gradient function in Python, save the resulting XLA HLO, then quickly evaluate both via the C++ interface.

(I'm gonna close this issue because I don't think there's any concrete work to do here, but feel free to continue commenting.)

[dawgster]

Hi @skye!
Could you give a few pointers as to how to get started on this? Ive went through the C++ client you linked but havnt found a possibility to load an HLO file at first or second glances. Ive already built a network in jax, and transformed it to HLO. The final piece of the puzzle is to run it via that interface. Any help would be greatly appreciated. Thanks!

[hawkinsp]

@dawgster I think the key thing to notice is that the XlaComputation passed to LocalClient::Compile() has a constructor that takes an HloModuleProto. Does that help?

[dawgster]

@hawkinsp Yes, that helped quite a bit. Thanks for your help!

[zhangqiaorjc]

#5337 shows an explicit example

It can be adapted easily for GPU client.

[uduse]

Do you think it's a good alternative to call JAX python code in C++?

[Roy-Kid]

Do you think it's a good alternative to call JAX python code in C++?

something like #include<Python.h>? I wonder if there is a more elegant approach to invoke jax in c++. I need to write a plugin for OpenMM and the reason I use jax is that it is small and flexible than tf.