[ci] Fix bug in test_ndarray.py when using the latest version of JAX
#1375
| Job | Run time |
|---|---|
| 2m 13s | |
| 1m 2s | |
| 1m 46s | |
| 47s | |
| 58s | |
| 1m 37s | |
| 46s | |
| 43s | |
| 1m 40s | |
| 44s | |
| 49s | |
| 57s | |
| 14m 2s |