jax

31bdc301 - moved all to the standard custom_call.cc instead of customized file