Collect recursively and filter GPU tests using jax_test_gpu tag

[andportnoy]

No description provided.

[andportnoy]

Needs jax-ml/jax#24218.

[yhtang]

It looks like JAX unit tests resulted in 20k+ errors and my browser just spent the last 10 mins loading the error log...

[andportnoy]

Yeah this confuses me because the errors that I see are all of the form

RuntimeError: Backend 'tpu' failed to initialize: INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory. Available backends are ['cpu', 'cuda']

in tests like //tests:blocked_sampler_test_tpu. But one would think these tests shouldn't be collected when filtering with --test_tag_filters=jax_test_gpu"...

[andportnoy]

The full command is

bazel test \
  --//jax:build_jaxlib=false \
  //tests/... \
  --test_tag_filters=jax_test_gpu \
  --@local_config_cuda//:enable_cuda \
  --cache_test_results=no \
  --test_timeout=600 \
  --test_tag_filters=-multiaccelerator \
  --test_env=JAX_SKIP_SLOW_TESTS=1 \
  --test_env=JAX_ACCELERATOR_COUNT=8 \
  --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform \
  --test_env=XLA_PYTHON_CLIENT_PREALLOCATE=false \
  --test_output=errors \
  --java_runtime_version=remotejdk_11 \
  --run_under /opt/jax/build/parallel_accelerator_execute.sh \
  --local_test_jobs=64 \
  --test_env=JAX_TESTS_PER_ACCELERATOR=8 \
  --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow

I wonder if --test_tag_filters cannot be repeated and so --test_tag_filters=-multiaccelerator overrides --test_tag_filters=jax_test_gpu?

Nope, --test_tag_filters cannot be repeated: bazelbuild/bazel#7322.

[andportnoy]

jax-ml/jax#24322 fixes //tests/mosaic:flash_attention_gpu on V100/A100.

[olupton]

LGTM
//tests/mosaic:flash_attention_gpu is failing, but I don't think that should block merging this

[andportnoy]

//tests/mosaic:flash_attention_gpu is failing

Fixed in jax-ml/jax#24467.