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]) – 組譯碼產生的結果的形狀和 dtype。

返回:

組譯碼產生的結果。

返回類型:

Sequence[jax.Array]