add test for raise

apache · Feb 12, 2021 · 7ede696 · 7ede696
1 parent a094577
commit 7ede696
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py
@@ -1133,3 +1133,12 @@ def test_onnx_export_take(tmp_path, dtype, axis, mode):
     op_export_test('take1', M1, [x, y], tmp_path)
     M2 = def_model('take', axis=axis, mode=mode)
     op_export_test('take2', M2, [x, y], tmp_path)
+
+
+@pytest.mark.parametrize('dtype', ['int32', 'int64', 'float16', 'float32', 'float64'])
+@pytest.mark.parametrize('axis', [-3, -2, -1, 0, 1, 2])
+def test_onnx_export_take_raise(tmp_path, dtype, axis):
+    x = mx.nd.random.normal(0, 10, (3, 4, 5)).astype(dtype)
+    y = mx.random.randint(0, 4, (6, 7)).astype(dtype)
+    M = def_model('take', axis=axis, mode='raise')
+    op_export_test('take', M, [x, y], tmp_path)