jax.experimental.pallas.triton.elementwise_inline_asm#

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

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

参数:

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

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

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

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

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

返回:

汇编代码产生的结果。

返回类型:

Sequence[jax.Array]