more test for aten.linear

pytorch · Feb 3, 2025 · 10c124d · 10c124d
1 parent 1327602
commit 10c124d
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/torchax/test/test_core_aten_ops.py b/torchax/test/test_core_aten_ops.py
@@ -4490,15 +4490,25 @@ def test_aten_einsum(self):
                            rtol=1e-2, check_dtype=True)
 
   def test_aten_linear(self):
+    # with bias
     args = (
-        torch.randn((2, 2), dtype=torch.float16),
-        torch.randn((2, 2), dtype=torch.float16),
+        torch.randn((2, 4), dtype=torch.float16),
+        torch.randn((2, 4), dtype=torch.float16),
         torch.randn((2, ), dtype=torch.float16),
     )
     kwargs = dict()
     run_export_and_compare(self, torch.ops.aten.linear, args, kwargs, atol=1e-2,
                            rtol=1e-2, check_dtype=True)
 
+    # without bias
+    args = (
+        torch.randn((2, 4), dtype=torch.float16),
+        torch.randn((2, 4), dtype=torch.float16),
+    )
+    kwargs = dict()
+    run_export_and_compare(self, torch.ops.aten.linear, args, kwargs, atol=1e-2,
+                           rtol=1e-2, check_dtype=True)
+
 
 if __name__ == "__main__":
   test_base.main()