triton.language.xor_sum

triton.language.xor_sum = <function xor_sum>

Returns the xor sum of all elements in the input tensor along the provided axis

参数:

• input (Tensor) -- the input values

• axis (int) -- the dimension along which the reduction should be done. If None, reduce all dimensions

• keep_dims (bool) -- if true, keep the reduced dimensions with length 1

This function can also be called as a member function on tensor, as x.xor_sum(...) instead of xor_sum(x, ...). .. rubric:: Example

import torch
import torch_npu
import triton
import triton.language as tl


@triton.jit
def triton_xorsum_2d(in_ptr0, out_ptr0, dim: tl.constexpr, M: tl.constexpr, N: tl.constexpr, MNUMEL: tl.constexpr,
                     NNUMEL: tl.constexpr):
    mblk_idx = tl.arange(0, MNUMEL)
    nblk_idx = tl.arange(0, NNUMEL)
    mmask = mblk_idx < M
    nmask = nblk_idx < N
    mask = (mmask[:, None]) & (nmask[None, :])
    idx = mblk_idx[:, None] * N + nblk_idx[None, :]
    x = tl.load(in_ptr0 + idx, mask=mask, other=0)
    tmp4 = tl.xor_sum(x, axis=dim)
    if dim == 0:
        tl.store(out_ptr0 + tl.arange(0, N), tmp4, None)
    else:
        tl.store(out_ptr0 + tl.arange(0, M), tmp4, None)


def test_xor_sum():
    M, N, dim = 4, 8, 1
    x = torch.randint(0, 128, (M, N), dtype=torch.int32).npu()
    out = torch.empty(M, dtype=torch.int32).npu()
    triton_xorsum_2d[1, 1, 1](x, out, dim, M, N, M, N)
    ref = torch.zeros(M, dtype=torch.int32)
    for i in range(M):
        val = 0
        for j in range(N):
            val ^= x[i, j].item()
        ref[i] = val
    assert torch.equal(out.cpu(), ref), "xor_sum result mismatch"
    print("xor_sum result:", out)


if __name__ == "__main__":
    test_xor_sum()

Special Restrictions

• DataType: Ascend A2/A3 does not support uint16, uint32, uint64; Ascend 950 does not fp8e4(E4M3), fp8e5(E5M2).

• keep_dims=True requires more test coverage; currently verified for 3D tensor with dim=2.