jax

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