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[, ...])

註冊自訂呼叫目標。