Jieying Luo 88685d8de0 Support bazel test without bazel build for CUDA PJRT plugin. ...

- Add build target for jax_plugins/ and jax_plugins/cuda for bazel test.
- Update jax_plugins/cuda/__init__.py to fallback to local `.so` file path.
- Add a flag --//jax:build_cuda_plugin to control whether to link in local cuda plugin.

The following command will test with cuda plugin:
```
bazel test tests:python_callback_test_gpu --test_output=all --test_filter=PythonCallbackTest.test_send_zero_dim_arrays_pure --config=tensorflow_testing_rbe_linux --config=rbe_linux_cuda12.2_nvcc_py3.9 --//jax:build_cuda_plugin=false
```

Default behavior (without `--//jax:build_cuda_plugin=false`) remains unchanged.

PiperOrigin-RevId: 582728477

2023-11-15 10:38:19 -08:00

..

__init__.py

Support bazel test without bazel build for CUDA PJRT plugin.

2023-11-15 10:38:19 -08:00

BUILD.bazel

Support bazel test without bazel build for CUDA PJRT plugin.

2023-11-15 10:38:19 -08:00

plugin_pyproject.toml

[PJRT C API] Rename the folder "plugins" to "jax_plugins".

2023-11-14 13:56:13 -08:00

plugin_setup.py

[PJRT C API] Rename the folder "plugins" to "jax_plugins".

2023-11-14 13:56:13 -08:00

pyproject.toml

[PJRT C API] Rename the folder "plugins" to "jax_plugins".

2023-11-14 13:56:13 -08:00

setup.py

[PJRT C API] Rename the folder "plugins" to "jax_plugins".

2023-11-14 13:56:13 -08:00