jax

e8888065 - Fix the pjit flakey test. The test was weirdly written in the first place. The current suspicion is that `x.copy()` inside jit made it flakey. Also delete some duplicate tests from the time of migrating to jax.Array