Clarifying whether jax.jit only works for Jax array and not numpy arrays

[prrathi]

Hi I was reading the following snippet https://github.com/google/jax/blob/0b87bf48f97ace10c7aee19c8f980788891a2df7/docs/sharded-computation.md?plain=1#L243-L262 in the documentation and saw the example function was wrapped in jax.jit but took in a numpy array. I had thought jit only works as intended with Jax arrays so wanted to clarify this. I ask this as a beginner and for my own understanding, not to nitpick the documentation!

[jakevdp]

Hi - thanks for the question! Functions wrapped in jax.jit only operate on JAX arrays, not on NumPy arrays. But for historical reasons, JIT-compiled functions accept NumPy arrays as inputs, and silently convert them to JAX arrays.