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]