[JAX-CFD] tweak installation & tests

The tweaks to the test suite ensure everything passes on my local
installation on my MacBook, and shave a few minutes off the execution
time.

If the README instructions are followed, the whole test-suite runs in about
two and a half minutes on my Mac. It excludes the "validation tests", which
are too slow to run on CPU.

PiperOrigin-RevId: 374442976