jax.lib 模块

此 jax.lib 包是一组内部工具和类型，用于连接 JAX 的 Python 前端和 XLA 后端。

jax.lib.xla_bridge

get_backend([platform])
get_compile_options(num_replicas, num_partitions)	返回要使用的编译选项，根据标志值推导得出。

jax.lib.xla_client

register_custom_call_target(name, fn[, ...])

注册一个自定义调用目标。