jax.process_indices

jax.process_indices(backend=None)

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

参数：

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

返回：

整数进程索引的列表。

返回类型：

list[int]