jax
db11efab - Migrate jaxlib to use a single common .so file for all C++ dependencies.

Commit
330 days ago
Migrate jaxlib to use a single common .so file for all C++ dependencies. The idea is to move all of the jaxlib contents into a single .so file, and have all of the other Python extensions be tiny stubs that reexport part of the larger .so file. This has two main benefits: * it reduces the size of the jaxlib wheel, by about 70-80MB when installed. The benefit of the change is that it avoid duplication between the MLIR CAPI code and the copy of MLIR in XLA. * it gives us flexibility to split and merge Python extensions as we see fit. Issue https://github.com/jax-ml/jax/issues/11225 PiperOrigin-RevId: 744855997
Author
Parents
Loading