triton.language.randn

triton.language.randn(seed, offset, n_rounds: constexpr = constexpr[10])

Given a seed scalar and an offset block, returns a block of random float32 in \(\mathcal{N}(0, 1)\).

Parameters:

• seed – The seed for generating random numbers.

• offsets – The offsets to generate random numbers for.

Example

import triton
import triton.language as tl
import torch
import math


@triton.jit
def kernel_randn(x_ptr, n_rounds: tl.constexpr, N: tl.constexpr, XBLOCK: tl.constexpr):
    block_offset = tl.program_id(0) * XBLOCK
    offsets = block_offset + tl.arange(0, XBLOCK)  # Block-level offset tensor
    mask = offsets < N
    rand_vals = tl.randn(5, 10 + offsets, n_rounds)  # Generate a block of random numbers at once
    tl.store(x_ptr + offsets, rand_vals, mask=mask)


def test_randn():
    shape = (1, 3)

    y_calf = torch.zeros(shape, dtype=eval('torch.float32')).npu()
    numel = y_calf.numel()
    ncore = 1 if numel < 32 else 32
    xblock = math.ceil(numel / ncore)
    kernel_randn[ncore, 1, 1](y_calf, 10, numel, xblock)


if __name__ == "__main__":
    test_randn()

DataType Support

平台	uint8	int8	uint16	int16	uint32	int32	uint64	int64	fp16	fp32	fp64	bf16	fp8e(e4m3)	fp8e5(e5m2)	bool
Ascend A2/A3	√	√	√	√	√	√	√	√	×	×	×	×	×	×	√
Ascend 950	√	√	√	√	√	√	√	√	×	×	×	×	×	×	√