Triton 是一種用於並行編程的語言和編譯器。它旨在提供一個基於 Python 的編程環境，以高效編寫自定義 DNN 計算內核，並能夠在現代 GPU 硬件上以最大吞吐量運行。

更多 Triton 中文文檔可訪問 →https://triton.hyper.ai/

triton.language.expand_dims(input, axis)

通過插入新的長度為 1 的維度來擴展張量的形狀。

軸索引是相對於生成的張量而言的，因此對於每個軸，result.shape[axis] 將為 1。

參數：

• input (tl.tensor) - 輸入張量。
• axis (int | Sequence[int] ) - 要添加新軸的索引。

該函數也可作為 tensor 的成員函數調用，使用 x.expand_dims(...) 而不是 expand_dims(x, ...)。