[Feature] Add a function to convert bitmasks to bool masks for development usage. (#426)

This PR provides a function to convert a bitmask to bool mask, the API
is:
```
def _bitmask_to_bool_mask(bit_mask: torch.Tensor, vocab_size: Optional[int] = None) -> torch.Tensor:
```

---------

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

Author: Seven-Streams

python/xgrammar/testing.py

@@ -262,6 +262,42 @@ 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:
+    """
+    Convert a bitmask tensor to a boolean mask tensor.
+
+    Parameters
+    ----------
+    bit_mask : torch.Tensor
+        The bitmask tensor to convert. Should be on CPU and of type int32.
+    vocab_size : Optional[int], default: None
+        The size of the vocabulary. If provided, the output mask will be cut to this size.
+
+    Returns
+    -------
+    bool_mask : torch.Tensor
+        The converted boolean mask tensor.
+    """
+
+    # Validate input.
+    if bit_mask.device.type != "cpu":
+        raise ValueError("bit_mask should be on CPU.")
+    if bit_mask.dtype != bitmask_dtype:
+        raise ValueError("bit_mask should be of type torch.int32.")
+
+    if vocab_size is None:
+        vocab_size = bit_mask.shape[1] * 32
+    if vocab_size > bit_mask.shape[1] * 32:
+        raise ValueError(
+            "vocab_size should be less than or equal to the size represented by bit_mask."
+        )
+
+    bool_mask = torch.zeros((bit_mask.shape[0], vocab_size), dtype=torch.bool)
+    for i in range(vocab_size):
+        bool_mask[:, i] = (bit_mask[:, i // 32] & (1 << (i % 32))) != 0
+    return bool_mask
+
+
 def _get_matcher_from_grammar_and_tokenizer_info(
     grammar: Union[Grammar, str], tokenizer_info: Optional[TokenizerInfo] = None, **kwargs
 ) -> GrammarMatcher:

tests/python/test_token_bitmask_operations.py

@@ -9,6 +9,7 @@
 
 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,
@@ -357,5 +358,36 @@ def test_apply_token_bitmask_inplace_indices(
         torch.testing.assert_close(logits, logits_expected)
 
 
+def test_bitmask_to_boolmask():
+    # 0xFFFF0000, 0x0000FFFF
+    bitmask = torch.tensor([[-65536, 65535]], dtype=torch.int32)
+    expected = torch.tensor(
+        [[False] * 16, [True] * 16, [True] * 16, [False] * 16], dtype=torch.bool
+    ).reshape(1, -1)
+    bool_mask = _bitmask_to_bool_mask(bitmask)
+    assert torch.equal(bool_mask, expected)
+
+    bool_mask_50 = _bitmask_to_bool_mask(bitmask, vocab_size=50)
+    expected_50 = expected[:, :50]
+    assert torch.equal(bool_mask_50, expected_50)
+
+
+batch__size__vocab__size = [
+    (4, 1000),
+    (1, 1024),
+    (16, 1024),
+    # not a multiple of 16.
+    (3, 817),
+]
+
+
+@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)
+    assert torch.equal(bool_mask, bool_mask_converted)
+
+
 if __name__ == "__main__":
     pytest.main(sys.argv)