jax.extend.backend 模块

backends()
backend_xla_version([platform])	返回后端的 XLA 版本。
clear_backends()	清除所有后端客户端，以便稍后创建新的后端客户端。
get_backend([platform])
get_compile_options(num_replicas, num_partitions)	返回根据标志值导出的编译选项。
get_default_device()
ifrt_proxy
register_backend_cache(cache, for_what)	将缓存注册到 JAX 的缓存管理系统中。
register_backend_factory(name, factory, *[, ...])