test

tests/test_special_ops.py

@@ -18,6 +18,7 @@
     UPSAMPLE_SHAPES,
     UT_SHAPES_1D,
     UT_SHAPES_2D,
+    KRON_SHAPES,
     gems_assert_close,
     gems_assert_equal,
     to_reference,
@@ -974,3 +975,20 @@ def test_sort(batch_size, hiddensize, descending, dtype, dim):
 
     gems_assert_close(res_value, ref_value, dtype)
     gems_assert_equal(res_index, ref_index)
+
+@pytest.mark.kron
+@pytest.mark.parametrize("shape", KRON_SHAPES)
+@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
+def test_accuracy_kron(shape, dtype):
+    inp1 = torch.randn(shape[0], dtype=dtype, device=flag_gems.device)
+    inp2 = torch.randn(shape[0], dtype=dtype, device=flag_gems.device)
+
+    ref_inp1 = to_reference(inp1)
+    ref_inp2 = to_reference(inp2)
+
+    ref_out = torch.kron(ref_inp1, ref_inp2)
+    with flag_gems.use_gems():
+        res_out = torch.kron(inp1, inp2)
+
+    gems_assert_equal(res_out, ref_out)
+