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]