FlagGems/tests/test_nonzero_numpy.py at 5b2ccd90379d5a2dde7a432b997a8dace5e3c6cc · Schopenhauer-loves-Hegel/FlagGems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import pytest
import torch

import flag_gems

from .accuracy_utils import (
    BOOL_TYPES,
    FLOAT_DTYPES,
    INT_DTYPES,
    REDUCTION_SHAPES,
    gems_assert_equal,
    to_reference,
)

NONZERO_SHAPES = REDUCTION_SHAPES


@pytest.mark.nonzero_numpy
@pytest.mark.parametrize("shape", NONZERO_SHAPES)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + INT_DTYPES + BOOL_TYPES)
def test_nonzero_numpy(shape, dtype):
    if dtype == torch.bool:
        inp = torch.randint(0, 2, shape, dtype=torch.int, device=flag_gems.device).to(
            dtype
        )
    elif dtype in INT_DTYPES:
        inp = torch.randint(-3, 3, shape, device=flag_gems.device).to(dtype)
    else:
        inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    ref_inp = to_reference(inp, False)

    ref_out = torch.ops.aten.nonzero_numpy(ref_inp)
    with flag_gems.use_gems():
        res_out = torch.ops.aten.nonzero_numpy(inp)

    assert len(res_out) == len(ref_out), "Number of output tensors should match"
    for res_t, ref_t in zip(res_out, ref_out):
        gems_assert_equal(res_t, ref_t)