triton.language.zeros

triton.language.zeros = <function zeros>

Returns a tensor filled with the scalar value 0 for the given shape and dtype.

参数:

• shape (tuple of ints) -- Shape of the new array, e.g., (8, 16) or (8, )

• dtype (DType) -- Data-type of the new array, e.g., tl.float16

Example

import torch
import triton
import triton.language as tl


@triton.jit
def zeros_2d_kernel(output_ptr, XB: tl.constexpr, YB: tl.constexpr):
    """
    Generate a 2D zero tensor with shape (XB, YB) and store it to output_ptr in row-major order.
    """
    xidx = tl.arange(0, XB)
    yidx = tl.arange(0, YB)

    ret = tl.zeros((XB, YB), dtype=tl.float32)

    oidx = xidx[:, None] * YB + yidx[None, :]

    tl.store(output_ptr + oidx, ret)


def test_zeros_2d():
    XB = 4
    YB = 8

    output = torch.zeros((XB, YB), dtype=torch.float32).npu()

    zeros_2d_kernel[(1, )](output, XB=XB, YB=YB)

    expected = torch.zeros((XB, YB), dtype=torch.float32)

    assert torch.allclose(output.cpu(), expected.cpu())


if __name__ == "__main__":
    test_zeros_2d()

Special Restrictions

• DataType: Ascend A2/A3 does not support uint16, uint32, uint64 (hardware limitation).