PythonCall With Jax: Fast inference w/ numpy but does not work with jax.grad need jax.numpy, which is slow, for gradient

[AlexLewandowski]

I need Jax for something that Zygote cannot do well (meta-learning) and someone recommended PythonCall as a solution to some issues I was having with PyCall.

So far, PythonCall has been great. Things work and it is generally quite quick.

There is one pain point: jax.grad does not work with numpy.array

[AlexLewandowski]

Apologies, seems like a draft was posted by mistake.