jax.process_indices

jax.process_indices(backend=None)

返回与后端关联的所有 JAX 进程索引的列表。

参数:

后端 (str | xla_client.Client | None) – 这是一个实验性功能，API 可能会发生变化。可选参数，一个字符串，表示 XLA 后端：'cpu'、'gpu' 或 'tpu'。

返回:

整数进程索引列表。

返回类型:

list[int]