jax.experimental.pallas.triton.elementwise_inline_asm#

jax.experimental.pallas.triton.elementwise_inline_asm(asm, *, args, constraints, pack, result_shape_dtypes)[源代码]#

内联汇编，应用逐元素操作。

参数:

  • asm (str) – 要运行的汇编代码。

  • args (Sequence[jax.Array]) – 传递给汇编代码的参数。

  • constraints (str) – LLVM 内联汇编约束。

  • pack (int) – 单个汇编代码实例预期从每个参数中获取的元素数量。

  • result_shape_dtypes (Sequence[jax.ShapeDtypeStruct]) – 汇编代码产生的结果的形状和数据类型。

返回:

汇编代码产生的结果。

返回类型:

Sequence[jax.Array]