Skip to content

Commit 9d9c072

Browse files
taozhiweitaozhiwei
authored
Merge branch 'master' into myfeature2
2 parents 9d7432e + a8a4452 commit 9d9c072

11 files changed

Lines changed: 765 additions & 623 deletions

‎tests/test_all.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import accuracy_utils as utils
7+
from . import conftest as cfg
8+
9+
if cfg.QUICK_MODE:
10+
FLOAT_DTYPES = [torch.float32]
11+
DIMS_LIST = [1]
12+
KIND_KEEPDIM_DIMS_SHAPE = [("normal", True, 1, utils.REDUCTION_SHAPES[0])]
13+
else:
14+
FLOAT_DTYPES = utils.FLOAT_DTYPES
15+
DIMS_LIST = [0, 1, [0, 1], [1, 0]]
16+
KIND_KEEPDIM_DIMS_SHAPE = list(
17+
zip(
18+
["normal", "allTrue"] * 2,
19+
[True, False] * 2,
20+
DIMS_LIST,
21+
utils.REDUCTION_SHAPES + [(7, 4, 11, 1)],
22+
)
23+
)
24+
25+
26+
@pytest.mark.all
27+
@pytest.mark.parametrize("shape", utils.REDUCTION_SHAPES)
28+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + [torch.bool])
29+
@pytest.mark.parametrize("kind", ["normal", "allTrue"])
30+
def test_all(shape, dtype, kind):
31+
if kind == "allTrue":
32+
inp = torch.ones(shape, dtype=dtype, device=flag_gems.device)
33+
else:
34+
inp = torch.randint(0, 2, shape, dtype=dtype, device="cpu").to(flag_gems.device)
35+
ref_inp = utils.to_reference(inp)
36+
37+
ref_out = torch.all(ref_inp)
38+
with flag_gems.use_gems():
39+
res_out = torch.all(inp)
40+
41+
utils.gems_assert_equal(res_out, ref_out)
42+
43+
44+
@pytest.mark.all_dims
45+
@pytest.mark.skipif(
46+
utils.SkipVersion("torch", "<2.2"), reason="Skipping Pytorch version."
47+
)
48+
@pytest.mark.parametrize("kind, keepdim, dim, shape", KIND_KEEPDIM_DIMS_SHAPE)
49+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + [torch.bool])
50+
def test_all_dims(shape, dim, keepdim, dtype, kind):
51+
if kind == "allTrue":
52+
inp = torch.ones(shape, dtype=dtype, device=flag_gems.device)
53+
else:
54+
inp = torch.randint(0, 2, shape, dtype=dtype, device="cpu").to(flag_gems.device)
55+
ref_inp = utils.to_reference(inp)
56+
57+
ref_out = torch.all(ref_inp, dim=dim, keepdim=keepdim)
58+
with flag_gems.use_gems():
59+
res_out = torch.all(inp, dim=dim, keepdim=keepdim)
60+
61+
utils.gems_assert_equal(res_out, ref_out)

‎tests/test_allclose.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import accuracy_utils as utils
7+
from . import conftest as cfg
8+
9+
if cfg.QUICK_MODE:
10+
FLOAT_DTYPES = [torch.float32]
11+
else:
12+
FLOAT_DTYPES = utils.FLOAT_DTYPES
13+
14+
15+
@pytest.mark.allclose
16+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
17+
@pytest.mark.parametrize("dtype", utils.ALL_FLOAT_DTYPES + utils.ALL_INT_DTYPES)
18+
@pytest.mark.parametrize("equal_nan", [False, True])
19+
@pytest.mark.parametrize("gen_nan", [0, 1, 2, 3, 4])
20+
def test_allclose(shape, dtype, equal_nan, gen_nan):
21+
# [gen_nan] 1: nan, 2: inf, 3: -inf, 4: inf vs -inf
22+
rtol = torch.rand(1, dtype=torch.float32, device=flag_gems.device).item() * (
23+
0.0001 if dtype in [torch.bfloat16, torch.float16] else 0.01
24+
)
25+
if dtype in utils.ALL_FLOAT_DTYPES:
26+
atol = (
27+
torch.finfo(dtype).tiny
28+
* torch.randint(0, 4, (1,), device=flag_gems.device).item()
29+
)
30+
inp1 = torch.full(shape, 1.234, dtype=dtype, device=flag_gems.device)
31+
inp2 = torch.full(shape, 1.234, dtype=dtype, device=flag_gems.device)
32+
if gen_nan:
33+
nan_num = torch.full(
34+
(1,),
35+
float("nan" if gen_nan == 1 else "inf"),
36+
dtype=dtype,
37+
device=flag_gems.device,
38+
)
39+
# FIXME: Neg doesn't support double on torch_musa, so workaround temporarily.
40+
inp1.view(-1)[0] = (
41+
(-nan_num.cpu()).to(flag_gems.device) if gen_nan == 3 else nan_num
42+
)
43+
inp2.view(-1)[0] = (
44+
(-nan_num.cpu()).to(flag_gems.device) if gen_nan >= 3 else nan_num
45+
)
46+
else:
47+
atol = (
48+
torch.finfo(torch.float16).eps
49+
* torch.randint(0, 10, (1,), device=flag_gems.device).item()
50+
)
51+
inp1 = torch.randint(-1000, 1000, shape, device=flag_gems.device).to(dtype)
52+
inp2 = torch.randint(-1000, 1000, shape, device=flag_gems.device).to(dtype)
53+
54+
ref_inp1 = utils.to_reference(inp1, False)
55+
ref_inp2 = utils.to_reference(inp2, False)
56+
57+
with flag_gems.use_gems():
58+
res_out = torch.allclose(inp1, inp2, rtol, atol, equal_nan=equal_nan)
59+
60+
ref_out = torch.allclose(ref_inp1, ref_inp2, rtol, atol, equal_nan=equal_nan)
61+
62+
assert res_out == ref_out

‎tests/test_any.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import accuracy_utils as utils
7+
from . import conftest as cfg
8+
9+
if cfg.QUICK_MODE:
10+
FLOAT_DTYPES = [torch.float32]
11+
DIMS_LIST = [1]
12+
KIND_KEEPDIM_DIMS_SHAPE = [("normal", True, 1, utils.REDUCTION_SHAPES[0])]
13+
else:
14+
FLOAT_DTYPES = utils.FLOAT_DTYPES
15+
DIMS_LIST = [0, 1, [0, 1], [1, 0]]
16+
KIND_KEEPDIM_DIMS_SHAPE = list(
17+
zip(
18+
["normal", "allTrue"] * 2,
19+
[True, False] * 2,
20+
DIMS_LIST,
21+
utils.REDUCTION_SHAPES + [(7, 4, 11, 1)],
22+
)
23+
)
24+
25+
26+
@pytest.mark.any
27+
@pytest.mark.parametrize("shape", utils.REDUCTION_SHAPES)
28+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + [torch.bool])
29+
@pytest.mark.parametrize("kind", ["normal", "allFalse"])
30+
def test_any(shape, dtype, kind):
31+
if kind == "allFalse":
32+
inp = torch.zeros(shape, dtype=dtype, device=flag_gems.device)
33+
else:
34+
inp = torch.randint(0, 2, shape, dtype=dtype, device="cpu").to(flag_gems.device)
35+
ref_inp = utils.to_reference(inp)
36+
37+
ref_out = torch.any(ref_inp)
38+
with flag_gems.use_gems():
39+
res_out = torch.any(inp)
40+
41+
utils.gems_assert_equal(res_out, ref_out)
42+
43+
44+
@pytest.mark.any_dims
45+
@pytest.mark.skipif(
46+
utils.SkipVersion("torch", "<2.2"), reason="Skipping Pytorch version."
47+
)
48+
@pytest.mark.parametrize("kind, keepdim, dim, shape", KIND_KEEPDIM_DIMS_SHAPE)
49+
@pytest.mark.parametrize("dtype", FLOAT_DTYPES + [torch.bool])
50+
def test_any_dims(shape, dim, keepdim, dtype, kind):
51+
if kind == "allFalse":
52+
inp = torch.zeros(shape, dtype=dtype, device=flag_gems.device)
53+
else:
54+
inp = torch.randint(0, 2, shape, dtype=dtype, device="cpu").to(flag_gems.device)
55+
ref_inp = utils.to_reference(inp)
56+
57+
ref_out = torch.any(ref_inp, dim=dim, keepdim=keepdim)
58+
with flag_gems.use_gems():
59+
res_out = torch.any(inp, dim=dim, keepdim=keepdim)
60+
61+
utils.gems_assert_equal(res_out, ref_out)

0 commit comments

Comments
 (0)