[ci] Fix bug in test_ndarray.py when using the latest version of JAX
#1458
| Job | Run time |
|---|---|
| 25m 30s | |
| 22m 34s | |
| 14m 13s | |
| 21m 58s | |
| 14m 41s | |
| 28m 12s | |
| 22m 12s | |
| 25m 54s | |
| 13m 32s | |
| 9m 33s | |
| 12m 36s | |
| 25m 7s | |
| 3h 56m 2s |