Skip to content

Commit 23d3cf8

Browse files
tengqmtengqm
authored
Split tensor constructor accuracy tests (#2623)
1 parent 7dd58fe commit 23d3cf8

18 files changed

Lines changed: 597 additions & 451 deletions

‎tests/test_arange.py‎

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,30 +4,14 @@
44
import flag_gems
55

66
from . import accuracy_utils as utils
7-
from . import conftest as cfg
87

98

109
@pytest.mark.arange
11-
@pytest.mark.parametrize("start", utils.ARANGE_START)
12-
@pytest.mark.parametrize("step", [1, 2, 5])
13-
@pytest.mark.parametrize("end", [128, 256, 1024])
14-
@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.ALL_INT_DTYPES + [None])
15-
@pytest.mark.parametrize("device", [flag_gems.device, None])
16-
@pytest.mark.parametrize(
17-
"pin_memory", [False, None]
18-
) # Since triton only target to GPU, pin_memory only used in CPU tensors.
19-
def test_arange(start, step, end, dtype, device, pin_memory):
20-
ref_out = torch.arange(
21-
start,
22-
end,
23-
step,
24-
dtype=dtype,
25-
device="cpu" if cfg.TO_CPU else device,
26-
pin_memory=pin_memory,
27-
)
10+
@pytest.mark.parametrize("end", [10, 100, 1000, 5.0])
11+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.int64])
12+
def test_arange(end, dtype):
2813
with flag_gems.use_gems():
29-
res_out = torch.arange(
30-
start, end, step, dtype=dtype, device=device, pin_memory=pin_memory
31-
)
14+
res_out = torch.arange(end, dtype=dtype, device=flag_gems.device)
15+
ref_out = torch.arange(end, dtype=dtype, device="cpu")
3216

33-
utils.gems_assert_equal(res_out, ref_out)
17+
utils.gems_assert_equal(res_out.cpu(), ref_out)

‎tests/test_arange_start.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
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+
device = flag_gems.device
10+
11+
12+
@pytest.mark.arange_start
13+
@pytest.mark.parametrize("start", utils.ARANGE_START)
14+
@pytest.mark.parametrize("step", [1, 2, 5])
15+
@pytest.mark.parametrize("end", [128, 256, 1024])
16+
@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES + utils.ALL_INT_DTYPES + [None])
17+
@pytest.mark.parametrize("device", [flag_gems.device, None])
18+
# Since triton only target to GPU, pin_memory only used in CPU tensors.
19+
@pytest.mark.parametrize("pin_memory", [False, None])
20+
def test_arange(start, step, end, dtype, device, pin_memory):
21+
ref_out = torch.arange(
22+
start,
23+
end,
24+
step,
25+
dtype=dtype,
26+
device="cpu" if cfg.TO_CPU else device,
27+
pin_memory=pin_memory,
28+
)
29+
30+
with flag_gems.use_gems():
31+
res_out = torch.arange(
32+
start, end, step, dtype=dtype, device=device, pin_memory=pin_memory
33+
)
34+
35+
utils.gems_assert_equal(res_out, ref_out)

‎tests/test_eye.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import pytest
2+
import torch
3+
from packaging import version
4+
5+
import flag_gems
6+
7+
from . import accuracy_utils as utils
8+
from . import conftest as cfg
9+
10+
device = flag_gems.device
11+
12+
13+
@pytest.mark.eye
14+
@pytest.mark.parametrize(
15+
"shape",
16+
[
17+
(256, 1024),
18+
(1024, 256),
19+
(8192, 4096),
20+
(4096, 8192),
21+
]
22+
+ [(2**d, 2**d) for d in range(7, 13)],
23+
)
24+
@pytest.mark.parametrize(
25+
"dtype", utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES + utils.BOOL_TYPES
26+
)
27+
def test_eye(shape, dtype):
28+
if (
29+
cfg.TO_CPU
30+
and dtype == torch.bfloat16
31+
and version.parse(torch.__version__) < version.parse("2.5.0")
32+
):
33+
pytest.skip("BFloat16 not supported on CPU in torch<2.5.0")
34+
35+
n, m = shape
36+
37+
# test eye(n, m) without dtype
38+
with flag_gems.use_gems():
39+
res_out = torch.eye(n, m, device=flag_gems.device)
40+
41+
utils.gems_assert_equal(
42+
res_out, torch.eye(n, m, device="cpu" if cfg.TO_CPU else device)
43+
)
44+
45+
# with dtype
46+
with flag_gems.use_gems():
47+
res_out = torch.eye(n, m, dtype=dtype, device=flag_gems.device)
48+
49+
utils.gems_assert_equal(
50+
res_out,
51+
torch.eye(n, m, dtype=dtype, device="cpu" if cfg.TO_CPU else device),
52+
)
53+
54+
# test eye(n)
55+
with flag_gems.use_gems():
56+
res_out = torch.eye(n, device=flag_gems.device)
57+
58+
utils.gems_assert_equal(
59+
res_out, torch.eye(n, device="cpu" if cfg.TO_CPU else device)
60+
)
61+
62+
# with dtype
63+
with flag_gems.use_gems():
64+
res_out = torch.eye(n, dtype=dtype, device=flag_gems.device)
65+
66+
utils.gems_assert_equal(
67+
res_out,
68+
torch.eye(n, dtype=dtype, device="cpu" if cfg.TO_CPU else device),
69+
)

‎tests/test_full.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
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+
device = flag_gems.device
10+
11+
12+
@pytest.mark.full
13+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
14+
@pytest.mark.parametrize(
15+
"dtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
16+
)
17+
@pytest.mark.parametrize("fill_value", [3.1415926, 2, False])
18+
def test_full(shape, dtype, fill_value):
19+
# without dtype
20+
ref_out = torch.full(shape, fill_value, device="cpu" if cfg.TO_CPU else device)
21+
22+
with flag_gems.use_gems():
23+
res_out = torch.full(shape, fill_value, device=flag_gems.device)
24+
25+
utils.gems_assert_equal(res_out, ref_out)
26+
27+
# with dtype
28+
ref_out = torch.full(
29+
shape, fill_value, dtype=dtype, device="cpu" if cfg.TO_CPU else device
30+
)
31+
32+
with flag_gems.use_gems():
33+
res_out = torch.full(shape, fill_value, dtype=dtype, device=flag_gems.device)
34+
35+
utils.gems_assert_equal(res_out, ref_out)

‎tests/test_full_like.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import math
2+
3+
import pytest
4+
import torch
5+
6+
import flag_gems
7+
8+
from . import accuracy_utils as utils
9+
10+
device = flag_gems.device
11+
12+
13+
@pytest.mark.full_like
14+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
15+
@pytest.mark.parametrize(
16+
"dtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
17+
)
18+
@pytest.mark.parametrize(
19+
"xdtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
20+
)
21+
@pytest.mark.parametrize(
22+
"fill_value", [3.1415926, 2, False, float("inf"), float("nan")]
23+
)
24+
def test_full_like(shape, dtype, xdtype, fill_value):
25+
if isinstance(fill_value, float) and (
26+
math.isinf(fill_value) or math.isnan(fill_value)
27+
):
28+
if dtype not in utils.ALL_FLOAT_DTYPES:
29+
pytest.skip("Skipping inf/nan test for non-float dtypes")
30+
31+
inp = torch.empty(size=shape, dtype=dtype, device=device)
32+
ref_inp = utils.to_reference(inp)
33+
34+
# without dtype
35+
ref_out = torch.full_like(ref_inp, fill_value)
36+
with flag_gems.use_gems():
37+
res_out = torch.full_like(inp, fill_value)
38+
utils.gems_assert_equal(res_out, ref_out, equal_nan=True)
39+
40+
# with dtype
41+
ref_out = torch.full_like(ref_inp, fill_value, dtype=dtype)
42+
with flag_gems.use_gems():
43+
res_out = torch.full_like(inp, fill_value, dtype=dtype)
44+
45+
utils.gems_assert_equal(res_out, ref_out, equal_nan=True)

‎tests/test_new_full.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import math
2+
3+
import pytest
4+
import torch
5+
6+
import flag_gems
7+
8+
from . import accuracy_utils as utils
9+
10+
device = flag_gems.device
11+
12+
13+
@pytest.mark.new_full
14+
@pytest.mark.skip(reason="This test fails on line 50.")
15+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
16+
@pytest.mark.parametrize(
17+
"dtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
18+
)
19+
@pytest.mark.parametrize(
20+
"xdtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
21+
)
22+
@pytest.mark.parametrize(
23+
"fill_value", [3.1415926, 2, False, float("inf"), float("nan")]
24+
)
25+
def test_new_full(shape, dtype, xdtype, fill_value):
26+
inp = torch.empty(size=shape, dtype=dtype, device=device)
27+
ref_inp = utils.to_reference(inp)
28+
29+
# without dtype: output dtype inherits from self (dtype), skip if dtype doesn't support inf/nan
30+
if isinstance(fill_value, float) and (
31+
math.isinf(fill_value) or math.isnan(fill_value)
32+
):
33+
if dtype not in utils.ALL_FLOAT_DTYPES:
34+
pytest.skip("Skipping inf/nan test for non-float dtypes")
35+
36+
ref_out = ref_inp.new_full(shape, fill_value)
37+
with flag_gems.use_gems():
38+
res_out = inp.new_full(shape, fill_value)
39+
40+
utils.gems_assert_equal(res_out, ref_out, equal_nan=True)
41+
42+
# with dtype: output dtype is xdtype, skip if xdtype doesn't support inf/nan
43+
if isinstance(fill_value, float) and (
44+
math.isinf(fill_value) or math.isnan(fill_value)
45+
):
46+
if xdtype not in utils.ALL_FLOAT_DTYPES:
47+
pytest.skip("Skipping inf/nan test for non-float dtypes")
48+
49+
ref_out = ref_inp.new_full(shape, fill_value, dtype=xdtype)
50+
with flag_gems.use_gems():
51+
res_out = inp.new_full(shape, fill_value, dtype=xdtype)
52+
53+
utils.gems_assert_equal(res_out, ref_out, equal_nan=True)

‎tests/test_one_hot.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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+
device = flag_gems.device
10+
11+
12+
@pytest.mark.one_hot
13+
def test_one_hot():
14+
gems_one_hot = flag_gems.one_hot
15+
16+
dev_type = torch.device(device).type
17+
expected_device = "cpu" if cfg.TO_CPU else device
18+
19+
x = torch.tensor([3, 4, 1, 0], device=device, dtype=torch.int64)
20+
t = gems_one_hot(x)
21+
expected = torch.tensor(
22+
[[0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [0, 1, 0, 0, 0], [1, 0, 0, 0, 0]],
23+
device=expected_device,
24+
)
25+
utils.gems_assert_equal(t, expected)
26+
27+
t = gems_one_hot(x, -1)
28+
expected = torch.tensor(
29+
[[0, 0, 0, 1, 0], [0, 0, 0, 0, 1], [0, 1, 0, 0, 0], [1, 0, 0, 0, 0]],
30+
device=expected_device,
31+
)
32+
utils.gems_assert_equal(t, expected)
33+
34+
t = gems_one_hot(x, 6)
35+
expected = torch.tensor(
36+
[
37+
[0, 0, 0, 1, 0, 0],
38+
[0, 0, 0, 0, 1, 0],
39+
[0, 1, 0, 0, 0, 0],
40+
[1, 0, 0, 0, 0, 0],
41+
],
42+
device=expected_device,
43+
)
44+
utils.gems_assert_equal(t, expected)
45+
46+
x2 = torch.tensor([[3, 4], [1, 0]], device=device, dtype=torch.int64)
47+
t = gems_one_hot(x2)
48+
expected = torch.tensor(
49+
[[[0, 0, 0, 1, 0], [0, 0, 0, 0, 1]], [[0, 1, 0, 0, 0], [1, 0, 0, 0, 0]]],
50+
device=expected_device,
51+
)
52+
utils.gems_assert_equal(t, expected)
53+
54+
x0 = torch.tensor(4, device=device, dtype=torch.int64)
55+
t = gems_one_hot(x0)
56+
expected = torch.tensor([0, 0, 0, 0, 1], device=expected_device)
57+
utils.gems_assert_equal(t, expected)
58+
59+
x_empty = torch.empty([4, 0], dtype=torch.long, device=device)
60+
t = gems_one_hot(x_empty, 100)
61+
expected = torch.empty([4, 0, 100], dtype=torch.long, device=expected_device)
62+
utils.gems_assert_equal(t, expected)
63+
64+
if dev_type not in ("cuda", "xla", "mps"):
65+
bad = torch.tensor([3, 4, -1, 0], dtype=torch.long)
66+
with pytest.raises(RuntimeError):
67+
gems_one_hot(bad.to(device), -1)
68+
69+
bad = torch.tensor([3, 4, 1, 0], dtype=torch.long)
70+
with pytest.raises(RuntimeError):
71+
gems_one_hot(bad.to(device), 3)
72+
73+
with pytest.raises(RuntimeError):
74+
gems_one_hot(torch.empty([4, 0], dtype=torch.long, device=device))
75+
76+
with pytest.raises(RuntimeError):
77+
gems_one_hot(torch.tensor([3, 4, 1, 0], dtype=torch.long, device=device), -2)

‎tests/test_ones.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
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+
device = flag_gems.device
10+
11+
12+
@pytest.mark.ones
13+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
14+
@pytest.mark.parametrize(
15+
"dtype", utils.BOOL_TYPES + utils.ALL_INT_DTYPES + utils.ALL_FLOAT_DTYPES
16+
)
17+
def test_ones(shape, dtype):
18+
# without dtype
19+
with flag_gems.use_gems():
20+
res_out = torch.ones(shape, device=flag_gems.device)
21+
22+
utils.gems_assert_equal(
23+
res_out, torch.ones(shape, device="cpu" if cfg.TO_CPU else device)
24+
)
25+
26+
# with dtype
27+
with flag_gems.use_gems():
28+
res_out = torch.ones(shape, dtype=dtype, device=flag_gems.device)
29+
30+
utils.gems_assert_equal(
31+
res_out, torch.ones(shape, dtype=dtype, device="cpu" if cfg.TO_CPU else device)
32+
)

0 commit comments

Comments
 (0)