[ci] Fix bug in test_ndarray.py
when using the latest version of JAX
#1375
Loading
test_ndarray.py
when using the latest version of JAX
#1375