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