xla
xla copied to clipboard
openxla
A machine learning compiler for GPUs, CPUs, and ML accelerators
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Allow disabling nvshmem for build targets.
Cleanup unused visibility rules for XProf
Add JAX Windows presubmit job to XLA repo.
Update layout only when `allow_spmd_sharding_propagation_to_output/parameter` is true. Old GSPMD propagation needs them since they do not have the concept of open/closed sharding. In Shardy with sdy-round-trip, JAX creates the correct...
Include shape info to (Custom)KernelThunk::buffer_uses
Add a smoke test to PJRT wheel builds. It dlopens the shared object and checks that the PJRT API loads.