jax
b9c72967 - Add an interface for IFRT backends to specify executable serializability and use this information to automatically disable JAX compilation cache for backends that do not implement executable serialization

Commit
2 years ago
Add an interface for IFRT backends to specify executable serializability and use this information to automatically disable JAX compilation cache for backends that do not implement executable serialization This CL adds `supports_executable_serialization` as an IFRT client attribute indicating whether the executables produced by a given IFRT implementation are serializable or not. This is based on a principle where `xla::ifrt::Client::attributes()` returns a set of attributes representing the "capabilities" of an IFRT implementation, so that the users of IFRT can act based on such capabilities without having to know the exact backend that they are using. This change is backward compatible as IFRT backends that do not implement `supports_executable_serialization` are assumed to implement executable serialization. PiperOrigin-RevId: 606799188
Author
Committer
Parents
Loading