Improve test case for test_vmap_fix

EMI-Group · Jan 6, 2025 · 3e1d31e · 3e1d31e
1 parent 8d415de
commit 3e1d31e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/unit_test/core/test_vmap_fix.py b/unit_test/core/test_vmap_fix.py
@@ -51,6 +51,12 @@ def test_distance_fn_with_mask(self):
         self.assertIsNotNone(distances(self.costs[:-1], self.mask[:-1]))
 
     def test_distance_fn_without_mask(self):
+        distances = jit(
+            distance_fn, trace=True, lazy=False, example_inputs=(self.costs,)
+        )
+        self.assertIsNotNone(distances(self.costs))
+
+    def test_distance_fn_with_none(self):
         distances = jit(
             distance_fn, trace=True, lazy=False, example_inputs=(self.costs, None)
         )