jax.lax.platform_dependent#

jax.lax.platform_dependent(*args, default=None, **per_platform)[source]#

分阶段输出平台特定代码。

在 JAX 中，实际运行计算的平台确定得非常晚，例如，基于数据所在的位置。当使用 AOT 降低或序列化时，计算可能会在不同的机器上编译和执行，甚至在降低时不可用的平台上执行。这意味着使用 Python 条件编写平台相关代码是不安全的，例如，基于当前的默认 JAX 平台。相反，可以使用 platform_dependent

用法

def cpu_code(*args): ...
def tpu_code(*args): ...
def other_platforms_code(*args): ...
res = platform_dependent(*args, cpu=cpu_code, tpu=tpu_code,
                         default=other_platforms_code)

当分阶段输出的代码在 CPU 上执行时，这等效于 cpu_code(*args)，在 TPU 上等效于 tpu_code(*args)，在任何其他平台上等效于 other_platforms_code(*args)。与 Python 条件不同，所有备选项都会被追踪并分阶段输出到 Jaxpr。这类似于 switch()，并根据它来实现，从中继承了在转换下的行为。

与 switch() 不同，对执行内容的的选择会更早做出：在大多数情况下，在降低平台已知时在降低期间；在罕见的多平台降低和序列化的情况下，StableHLO 代码将包含实际平台上的条件。此条件在编译平台已知时，在编译之前及时解决。这意味着编译器实际上永远不会看到条件。

参数:

  • *args (Any) – 传递给每个分支的 JAX 数组。可以是 PyTrees。

  • **per_platform (Callable[..., _T]) – 用于不同平台的分支。这些分支是使用 *args 调用的 JAX 可调用对象。关键字是平台名称，例如 ‘cpu’、‘tpu’、‘cuda’、‘rocm’。

  • default (Callable[..., _T] | None | None) – 用于 per_platform 中未提及的平台的可选默认分支。如果没有 default，则当代码为 per_platform 中未提及的平台降低时，将发生错误。

返回:

值 per_platform[execution_platform](*args)。