PythonCall w/ Jax: Fast inference w/ numpy but only jax.numpy works for jax.grad

I am trying to use Jax in my Julia codebase for something that Zygote cannot do well (meta-learning). 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` should work with `numpy.array` (it does in Python) but it errors with PythonCall:

```
PyException( as an abstract array; it does not have a dtype attribute")>)
````

Instead, `jax.numpy.array` works with `jax.grad` but it is slower.

I assume this is due to the way PythonCall handles non-copying conversions under the hood.

Is there any way to reconcile these two and get `jax.grad` to work with `numpy.array` and the optimizations for non-copying arrays?

I've included an MWE:

```
using PythonCall
using Flux

jax = pyimport("jax")
jnp = pyimport("jax.numpy")
np = pyimport("numpy")
stax = pyimport("jax.example_libraries.stax")
optimizers = pyimport("jax.example_libraries.optimizers")
optax = pyimport("optax")
random = pyimport("jax.random")

rngkey = random.PRNGKey(123)

dense = stax.Dense
Relu = stax.Relu

in_shape = (-1, 1)

learner_init, learner_apply = stax.serial(dense(1))
out_shape, learner_params = learner_init(input_shape=in_shape, rng=rngkey)

x = transpose(rand(Float32, 1, 1))

opt_init, opt_update, get_params = optimizers.adam(step_size=1e-3)

opt_state = opt_init(learner_params)

p = get_params(opt_state)

function mse_loss(p,x)
y_hat = learner_apply(p,x)
return jnp.mean(optax.l2_loss(y_hat, x))
end

###
### For inference, both np.array and jnp.array work fine
###

learner_apply(p, np.array(x)); # Faster if you check with @btime
learner_apply(p, jnp.array(x)); # Slower

###
### For gradients, only jnp.array works:
###

try
jax.grad(mse_loss)(p, jnp.array(x))
println("jax.numpy.array works with jax.grad but is slow")
println()
catch e
println(e)
end

try
jax.grad(mse_loss)(p, np.array(x))
catch e
println(e)
println()
println("numpy.array does not work with jax.grad")
end
```

The CondaPkg.toml file:
```
[deps]
python = "3.10"

[pip.deps]
jax = ""
jaxlib = ""
optax = ""
```

Start by running the MWE with the listed CondaPkg.toml dependencies and compare jax.grad(mse_loss) with np.array(x) versus jnp.array(x). Trace PythonCall's non-copying array conversion and JAX's dtype or abstract-array handling. Done means determining whether the fast NumPy path can work with jax.grad, or clearly documenting the limitation and workaround.

索引モデルが issue の本文から書いたものです。