提交 cae7c1a5 authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者: Brandon T. Willard

Use pytest's parameterize for TestSubtensor.test_ellipsis

上级 fd4c5d91
......@@ -537,10 +537,9 @@ class TestSubtensor(utt.OptimizationTestMixin):
ret = f()
assert ret.shape == (1, 1, 4)
def test_ellipsis(self):
numpy_n = np.arange(24, dtype=self.dtype).reshape((2, 3, 4))
n = self.shared(numpy_n)
test_cases = [
@pytest.mark.parametrize(
"length, op_type_opt, slice_",
[
(0, Subtensor, np.index_exp[...]),
(1, Subtensor, np.index_exp[..., 1]),
(1, Subtensor, np.index_exp[1, ...]),
......@@ -554,9 +553,11 @@ class TestSubtensor(utt.OptimizationTestMixin):
AdvancedSubtensor,
np.index_exp[..., np.newaxis, [1, 2]],
),
]
for length, op_type_opt, slice_ in test_cases:
],
)
def test_ellipsis(self, length, op_type_opt, slice_):
numpy_n = np.arange(24, dtype=self.dtype).reshape((2, 3, 4))
n = self.shared(numpy_n)
numpy_tval = numpy_n[slice_]
t = n[slice_]
tval = self.eval_output_and_check(t, op_type=op_type_opt, length=length)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论