Change testcase to FD style

Author: littledgg

tests/batch_invariant/test_batch_invariance_op_mm.py

@@ -0,0 +1,58 @@
+# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/test_batch_invariance.py
+
+import unittest
+
+import paddle
+
+from custom_ops.batch_invariant_ops import set_batch_invariant_mode
+
+
+class TestBatchInvariantForMM(unittest.TestCase):
+    def setUp(self):
+        """
+        Initialize the test environment
+        """
+        device = "gpu" if paddle.is_compiled_with_cuda() else "cpu"
+        paddle.set_device(device)
+
+    def test_batch_invariance(self, B: int = 2048, D: int = 4096, dtype=paddle.float32):
+        a = paddle.linspace(-100, 100, B * D, dtype=dtype).reshape(B, D)
+        b = paddle.linspace(-100, 100, D * D, dtype=dtype).reshape(D, D)
+
+        # Method 1: Matrix-vector multiplication (batch size 1)
+        out1 = paddle.mm(a[:1], b)
+
+        # Method 2: Matrix-matrix multiplication, then slice (full batch)
+        out2 = paddle.mm(a, b)[:1]
+
+        # Check if results are identical
+        diff = (out1 - out2).abs().max()
+        return diff.item() == 0, diff
+
+    def run_iters(self, iters=10, ass=False):
+        for dtype in [paddle.float32, paddle.bfloat16]:
+            is_deterministic = True
+            difflist = []
+            for i in range(iters):
+                isd, df = self.test_batch_invariance(dtype=dtype)
+                is_deterministic = is_deterministic and isd
+                difflist.append(df)
+            print(
+                f"Batch Deterministic: {is_deterministic} run-to-run max/min/diff {max(difflist)}/{min(difflist)}/{max(difflist)-min(difflist)} for {dtype} in {iters} iterations"
+            )
+            if ass:
+                assert max(difflist) == 0
+
+    def test_case(self):
+        # Test with standard Paddle (likely to show differences)
+        print("Standard Paddle:")
+        with set_batch_invariant_mode(False):
+            self.run_iters(ass=False)
+        # Test with batch-invariant operations
+        print("\nBatch-Invariant Mode:")
+        with set_batch_invariant_mode(True):
+            self.run_iters(ass=True)
+
+
+if __name__ == "__main__":
+    unittest.main()