[Format] Rename two testing functions to expose them. (#434)

This PR renames two functions to expose them.

Signed-off-by: Yuchuan <[email protected]>

Author: Seven-Streams

python/xgrammar/testing.py

@@ -233,7 +233,7 @@ def _is_single_token_bitmask(
     )
 
 
-def _bool_mask_to_bitmask(bool_mask: torch.Tensor) -> torch.Tensor:
+def bool_mask_to_bitmask(bool_mask: torch.Tensor) -> torch.Tensor:
     """Get the bitmask from bool mask. If the bool mask does not align with the 32-bit block
     size, it will add extra 1 paddings.
 
@@ -262,7 +262,7 @@ def _bool_mask_to_bitmask(bool_mask: torch.Tensor) -> torch.Tensor:
     return bitmask.to(torch.int32)
 
 
-def _bitmask_to_bool_mask(bit_mask: torch.Tensor, vocab_size: Optional[int] = None) -> torch.Tensor:
+def bitmask_to_bool_mask(bit_mask: torch.Tensor, vocab_size: Optional[int] = None) -> torch.Tensor:
     """
     Convert a bitmask tensor to a boolean mask tensor.
 

tests/python/test_token_bitmask_operations.py

@@ -9,10 +9,10 @@
 
 import xgrammar as xgr
 from xgrammar.testing import (
-    _bitmask_to_bool_mask,
-    _bool_mask_to_bitmask,
     _get_masked_tokens_from_bitmask,
     _is_single_token_bitmask,
+    bitmask_to_bool_mask,
+    bool_mask_to_bitmask,
 )
 
 _is_cuda_available = torch.cuda.is_available()
@@ -38,7 +38,7 @@ def test_allocate_reset_token_bitmask():
 @pytest.mark.parametrize("index", (0, 1))
 def test_get_masked_tokens_from_bitmask(token_mask_size: int, index: int):
     bool_mask = torch.randint(0, 2, (2, token_mask_size), dtype=torch.bool)
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
     expected = torch.where(~bool_mask[index])[0].tolist()
     assert _get_masked_tokens_from_bitmask(bitmask, token_mask_size, index) == expected
 
@@ -50,13 +50,13 @@ def test_is_single_token_bitmask():
     token_id = 100
 
     bool_mask = torch.zeros(batch, vocab_size, dtype=torch.bool)
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
     assert _is_single_token_bitmask(bitmask, vocab_size, batch_index) == (False, -1)
     bool_mask[batch_index, token_id] = True
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
     assert _is_single_token_bitmask(bitmask, vocab_size, batch_index) == (True, token_id)
     bool_mask[batch_index, token_id + 1] = True
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
     assert _is_single_token_bitmask(bitmask, vocab_size, batch_index) == (False, -1)
 
 
@@ -229,7 +229,7 @@ def test_apply_token_bitmask_inplace_kernel_large(
             bool_mask.scatter_(1, masked_positions, False)
             assert (bool_mask.sum(dim=-1) + masked_cnt == vocab_size).all().item()
 
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
 
     batch_indices = torch.arange(0, batch_size, stride, dtype=torch.int32)
 
@@ -238,7 +238,7 @@ def test_apply_token_bitmask_inplace_kernel_large(
         logits_expected[batch_indices], ~bool_mask[batch_indices], float("-inf")
     )
 
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
     if impl in ["cuda", "triton", "torch_compile"]:
         logits_gpu = logits.to("cuda")
         bitmask_gpu = bitmask.to("cuda")
@@ -340,7 +340,7 @@ def test_apply_token_bitmask_inplace_indices(
 
     logits = torch.ones(logits_batch_size, vocab_size, dtype=torch.float32)
     bool_mask = torch.zeros(bitmask_batch_size, vocab_size, dtype=torch.bool)
-    bitmask = _bool_mask_to_bitmask(bool_mask)
+    bitmask = bool_mask_to_bitmask(bool_mask)
 
     logits_expected = logits.clone()
     logits_expected[indices] = torch.masked_fill(
@@ -364,10 +364,10 @@ def test_bitmask_to_boolmask():
     expected = torch.tensor(
         [[False] * 16, [True] * 16, [True] * 16, [False] * 16], dtype=torch.bool
     ).reshape(1, -1)
-    bool_mask = _bitmask_to_bool_mask(bitmask)
+    bool_mask = bitmask_to_bool_mask(bitmask)
     assert torch.equal(bool_mask, expected)
 
-    bool_mask_50 = _bitmask_to_bool_mask(bitmask, vocab_size=50)
+    bool_mask_50 = bitmask_to_bool_mask(bitmask, vocab_size=50)
     expected_50 = expected[:, :50]
     assert torch.equal(bool_mask_50, expected_50)
 
@@ -384,8 +384,8 @@ def test_bitmask_to_boolmask():
 @pytest.mark.parametrize("batch_size, vocab_size", batch__size__vocab__size)
 def test_bool_mask_bitmask_roundtrip(batch_size: int, vocab_size: int):
     bool_mask = torch.randint(0, 2, (batch_size, vocab_size), dtype=torch.bool)
-    bitmask = _bool_mask_to_bitmask(bool_mask)
-    bool_mask_converted = _bitmask_to_bool_mask(bitmask, vocab_size=vocab_size)
+    bitmask = bool_mask_to_bitmask(bool_mask)
+    bool_mask_converted = bitmask_to_bool_mask(bitmask, vocab_size=vocab_size)
     assert torch.equal(bool_mask, bool_mask_converted)