xla
xla copied to clipboard
[GPU] Expose caching DebugOptions to python
Adds python bindings for xla_gpu_kernel_cache_file, xla_gpu_enable_llvm_module_compilation_parallelism and xla_gpu_per_fusion_autotune_cache_dir.
We would like to add some convenience features to JAX which will enable all caches with one flag/option (will open PR for that soon). This change is necessary for that.
@skye @hawkinsp for PjRT API changes
Is it possible to test those?
Thanks for reviewing, added test to xla_client_test.
What can we do to get this PR merged? @cheshire @skye