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