Relax NumPy version check in tests.
Relax NumPy version check in tests.

Only look at the major, minor, and patch parts of the NumPy version. There might be non-integer parts of the version for development builds of NumPy.

google/jax,google/jax,google/jax,google/jax