Fix JaxArrayWrapper TypeError with JAX tracers during JIT#211

Chessing234 wants to merge 1 commit intogoogle-deepmind:maingoogle-deepmind/graphcast:mainfrom

Chessing234:fix/issue-203-jax-tracer-ufuncChessing234/graphcast:fix/issue-203-jax-tracer-ufuncCopy head branch name to clipboard