"""Unit tests for srt/sampling/sampling_batch_info.py — no server, no model loading.""" from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="stage-a-test-cpu") import unittest from unittest.mock import MagicMock, patch import torch from sglang.srt.sampling.sampling_batch_info import ( SamplingBatchInfo, merge_bias_tensor, ) from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.test.test_utils import CustomTestCase VOCAB_SIZE = 32 DEVICE = "cpu" # Helper: construct a minimal SamplingBatchInfo def _make_info(batch_size=2, **overrides): """Create a SamplingBatchInfo with sane defaults for testing.""" defaults = dict( temperatures=torch.ones(batch_size, 1), top_ps=torch.ones(batch_size), top_ks=torch.full((batch_size,), TOP_K_ALL, dtype=torch.int32), min_ps=torch.zeros(batch_size), is_all_greedy=False, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False, vocab_size=VOCAB_SIZE, device=DEVICE, penalizer_orchestrator=MagicMock(is_required=False), ) defaults.update(overrides) return SamplingBatchInfo(**defaults) class TestMergeBiasTensor(CustomTestCase): def test_both_none_returns_none(self): """Test that merging two None tensors returns None.""" result = merge_bias_tensor(None, None, 2, 3, DEVICE, 0.0) self.assertIsNone(result) def test_both_present_concatenates(self): """Test that two present tensors are concatenated along batch dim.""" lhs = torch.ones(2, VOCAB_SIZE) rhs = torch.zeros(3, VOCAB_SIZE) result = merge_bias_tensor(lhs, rhs, 2, 3, DEVICE, 0.0) self.assertEqual(result.shape, (5, VOCAB_SIZE)) self.assertEqual(result[0, 0].item(), 1.0) self.assertEqual(result[3, 0].item(), 0.0) def test_lhs_none_fills_default(self): """Test that missing lhs is filled with default value before concatenation.""" rhs = torch.ones(3, VOCAB_SIZE) result = merge_bias_tensor(None, rhs, 2, 3, DEVICE, 0.0) self.assertEqual(result.shape, (5, VOCAB_SIZE)) # First 2 rows filled with default (0.0) self.assertEqual(result[0, 0].item(), 0.0) # Last 3 rows from rhs self.assertEqual(result[2, 0].item(), 1.0) def test_rhs_none_fills_default(self): """Test that missing rhs is filled with default value before concatenation.""" lhs = torch.ones(2, VOCAB_SIZE) result = merge_bias_tensor(lhs, None, 2, 3, DEVICE, 0.0) self.assertEqual(result.shape, (5, VOCAB_SIZE)) self.assertEqual(result[0, 0].item(), 1.0) # Last 3 rows filled with default (0.0) self.assertEqual(result[3, 0].item(), 0.0) def test_custom_default_value(self): """Test that a custom default (-1.0) fills the missing lhs rows.""" rhs = torch.ones(1, VOCAB_SIZE) result = merge_bias_tensor(None, rhs, 2, 1, DEVICE, -1.0) self.assertEqual(result[0, 0].item(), -1.0) self.assertEqual(result[1, 0].item(), -1.0) self.assertEqual(result[2, 0].item(), 1.0) # SamplingBatchInfo.__len__ class TestSamplingBatchInfoLen(CustomTestCase): def test_len_matches_batch_size(self): """Test that __len__ returns batch size (number of temperature rows).""" info = _make_info(batch_size=5) self.assertEqual(len(info), 5) class TestMergeCustomLogitProcessor(CustomTestCase): def test_both_none_returns_none(self): """Test that merging two None processor dicts returns None.""" result = SamplingBatchInfo.merge_custom_logit_processor( None, None, 2, 3, DEVICE ) self.assertIsNone(result) def test_same_key_merges_masks(self): """Test that same processor key concatenates the boolean masks.""" proc = MagicMock() lhs = {42: (proc, torch.tensor([True, False]))} rhs = {42: (proc, torch.tensor([False, True, True]))} result = SamplingBatchInfo.merge_custom_logit_processor(lhs, rhs, 2, 3, DEVICE) self.assertIn(42, result) self.assertEqual(result[42][1].shape[0], 5) self.assertTrue(result[42][1][0].item()) # from lhs self.assertFalse(result[42][1][1].item()) # from lhs self.assertTrue(result[42][1][3].item()) # from rhs def test_disjoint_keys(self): """Test that disjoint processor keys are merged with zero-filled padding.""" proc_a = MagicMock() proc_b = MagicMock() lhs = {1: (proc_a, torch.tensor([True, False]))} rhs = {2: (proc_b, torch.tensor([True]))} result = SamplingBatchInfo.merge_custom_logit_processor(lhs, rhs, 2, 1, DEVICE) # Key 1: lhs mask [True, False] + zero-filled rhs [False] self.assertEqual(result[1][1].shape[0], 3) self.assertTrue(result[1][1][0].item()) self.assertFalse(result[1][1][2].item()) # Key 2: zero-filled lhs [False, False] + rhs mask [True] self.assertEqual(result[2][1].shape[0], 3) self.assertFalse(result[2][1][0].item()) self.assertTrue(result[2][1][2].item()) def test_lhs_none_rhs_present(self): """Test that None lhs is treated as empty dict and rhs mask is padded.""" proc = MagicMock() rhs = {10: (proc, torch.tensor([True]))} result = SamplingBatchInfo.merge_custom_logit_processor(None, rhs, 2, 1, DEVICE) self.assertIn(10, result) self.assertEqual(result[10][1].shape[0], 3) # apply_logits_bias class TestApplyLogitsBias(CustomTestCase): def test_applies_linear_penalties(self): """Test that pre-accumulated linear penalties are added to logits.""" info = _make_info(batch_size=1) info.acc_linear_penalties = torch.tensor([[-1.0] * VOCAB_SIZE]) logits = torch.zeros(1, VOCAB_SIZE) info.apply_logits_bias(logits) self.assertAlmostEqual(logits[0, 0].item(), -1.0, places=5) def test_applies_logit_bias(self): """Test that per-token logit_bias is added to logits.""" info = _make_info(batch_size=1) bias = torch.zeros(1, VOCAB_SIZE) bias[0, 5] = 10.0 info.logit_bias = bias logits = torch.zeros(1, VOCAB_SIZE) info.apply_logits_bias(logits) self.assertAlmostEqual(logits[0, 5].item(), 10.0, places=5) self.assertAlmostEqual(logits[0, 0].item(), 0.0, places=5) def test_applies_vocab_mask(self): """Test that vocab_mask triggers the apply_mask_func callback.""" info = _make_info(batch_size=1) info.vocab_mask = torch.ones(1, VOCAB_SIZE) info.apply_mask_func = MagicMock() logits = torch.zeros(1, VOCAB_SIZE) info.apply_logits_bias(logits) info.apply_mask_func.assert_called_once() def test_applies_penalizer_orchestrator(self): """Test that a required orchestrator's apply() is called on logits.""" orch = MagicMock(is_required=True) info = _make_info(batch_size=1, penalizer_orchestrator=orch) logits = torch.zeros(1, VOCAB_SIZE) info.apply_logits_bias(logits) orch.apply.assert_called_once_with(logits) def test_no_bias_no_change(self): """Test that logits stay unchanged when no bias sources are set.""" info = _make_info(batch_size=1) info.acc_linear_penalties = None info.logit_bias = None info.vocab_mask = None logits = torch.zeros(1, VOCAB_SIZE) original = logits.clone() info.apply_logits_bias(logits) self.assertTrue(torch.equal(logits, original)) # update_penalties class TestUpdatePenalties(CustomTestCase): def test_required_creates_penalties_tensor(self): """Test that update_penalties allocates a zero tensor and calls orchestrator.apply.""" orch = MagicMock(is_required=True) info = _make_info(batch_size=2, penalizer_orchestrator=orch) info.update_penalties() self.assertIsNotNone(info.acc_linear_penalties) self.assertEqual(info.acc_linear_penalties.shape, (2, VOCAB_SIZE)) orch.apply.assert_called_once() def test_not_required_sets_none(self): """Test that update_penalties sets acc_linear_penalties to None when not required.""" orch = MagicMock(is_required=False) info = _make_info(batch_size=2, penalizer_orchestrator=orch) info.update_penalties() self.assertIsNone(info.acc_linear_penalties) # update_regex_vocab_mask class TestUpdateRegexVocabMask(CustomTestCase): def test_no_grammars_clears_mask(self): """Test that None grammars clears both vocab_mask and apply_mask_func.""" info = _make_info(batch_size=1) info.grammars = None info.update_regex_vocab_mask() self.assertIsNone(info.vocab_mask) self.assertIsNone(info.apply_mask_func) def test_empty_grammars_clears_mask(self): """Test that empty grammars list clears vocab_mask.""" info = _make_info(batch_size=1) info.grammars = [] info.update_regex_vocab_mask() self.assertIsNone(info.vocab_mask) def test_with_grammars_allocates_and_fills(self): """Test that an active grammar gets allocate, fill, and move called.""" grammar = MagicMock() grammar.finished = False grammar.is_terminated.return_value = False grammar.allocate_vocab_mask.return_value = torch.zeros(1, VOCAB_SIZE) grammar.move_vocab_mask.return_value = torch.zeros(1, VOCAB_SIZE) info = _make_info(batch_size=1) info.grammars = [grammar] info.update_regex_vocab_mask() grammar.allocate_vocab_mask.assert_called_once() grammar.fill_vocab_mask.assert_called_once() grammar.move_vocab_mask.assert_called_once() def test_mixed_grammars_only_active_fills(self): """Test that finished, terminated, and None grammars are skipped.""" active = MagicMock() active.finished = False active.is_terminated.return_value = False active.allocate_vocab_mask.return_value = torch.zeros(3, VOCAB_SIZE) active.move_vocab_mask.return_value = torch.zeros(3, VOCAB_SIZE) finished = MagicMock() finished.finished = True terminated = MagicMock() terminated.finished = False terminated.is_terminated.return_value = True info = _make_info(batch_size=3) info.grammars = [active, finished, terminated] info.update_regex_vocab_mask() active.fill_vocab_mask.assert_called_once() finished.fill_vocab_mask.assert_not_called() terminated.fill_vocab_mask.assert_not_called() # filter_batch class TestFilterBatch(CustomTestCase): def test_filter_keeps_correct_indices(self): """Test that filter retains rows at indices 0 and 2, dropping index 1.""" info = _make_info(batch_size=3) info.temperatures = torch.tensor([[1.0], [2.0], [3.0]]) info.top_ps = torch.tensor([0.9, 0.8, 0.7]) info.top_ks = torch.tensor([10, 20, 30], dtype=torch.int32) info.min_ps = torch.tensor([0.0, 0.1, 0.2]) info.logit_bias = torch.ones(3, VOCAB_SIZE) keep = torch.tensor([0, 2]) info.filter_batch([0, 2], keep) self.assertEqual(len(info), 2) self.assertAlmostEqual(info.temperatures[0, 0].item(), 1.0) self.assertAlmostEqual(info.temperatures[1, 0].item(), 3.0) self.assertAlmostEqual(info.top_ps[1].item(), 0.7) # logit_bias should also be filtered self.assertEqual(info.logit_bias.shape, (2, VOCAB_SIZE)) def test_filter_with_custom_logit_processor(self): """Test that filter updates both custom_params list and processor mask.""" proc = MagicMock() info = _make_info(batch_size=3) info.has_custom_logit_processor = True info.custom_logit_processor = {42: (proc, torch.tensor([True, False, True]))} info.custom_params = [{"a": 1}, {"b": 2}, {"c": 3}] keep = torch.tensor([0, 2]) info.filter_batch([0, 2], keep) self.assertEqual(info.custom_params, [{"a": 1}, {"c": 3}]) mask = info.custom_logit_processor[42][1] self.assertEqual(mask.shape[0], 2) def test_filter_removes_all_custom_processors(self): """Test cleanup when filter removes all requests using a processor.""" proc = MagicMock() info = _make_info(batch_size=3) info.has_custom_logit_processor = True info.custom_logit_processor = {42: (proc, torch.tensor([False, True, False]))} info.custom_params = [None, {"x": 1}, None] # Keep only index 0 and 2 — processor 42's mask becomes [False, False] keep = torch.tensor([0, 2]) info.filter_batch([0, 2], keep) self.assertFalse(info.has_custom_logit_processor) self.assertIsNone(info.custom_logit_processor) def test_filter_with_none_sampling_seed(self): """Test that filter preserves None sampling_seed without error.""" info = _make_info(batch_size=3) info.sampling_seed = None keep = torch.tensor([1]) info.filter_batch([1], keep) self.assertIsNone(info.sampling_seed) # merge_batch class TestMergeBatch(CustomTestCase): def test_merge_concatenates_tensors(self): """Test that merge concatenates temperature tensors from both batches.""" info1 = _make_info(batch_size=2) info1.temperatures = torch.tensor([[1.0], [2.0]]) info2 = _make_info(batch_size=1) info2.temperatures = torch.tensor([[3.0]]) info1.merge_batch(info2) self.assertEqual(len(info1), 3) self.assertAlmostEqual(info1.temperatures[2, 0].item(), 3.0) def test_merge_combines_flags(self): """Test that merge ANDs is_all_greedy and ORs need_*_sampling flags.""" info1 = _make_info( is_all_greedy=True, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False, ) info2 = _make_info( is_all_greedy=False, need_top_p_sampling=True, need_top_k_sampling=True, need_min_p_sampling=True, ) info1.merge_batch(info2) self.assertFalse(info1.is_all_greedy) # AND semantics self.assertTrue(info1.need_top_p_sampling) # OR semantics self.assertTrue(info1.need_top_k_sampling) # OR semantics self.assertTrue(info1.need_min_p_sampling) # OR semantics def test_merge_with_logit_bias(self): """Test that merge pads missing logit_bias with zeros before concatenation.""" info1 = _make_info(batch_size=1) info1.logit_bias = torch.ones(1, VOCAB_SIZE) info2 = _make_info(batch_size=1) info2.logit_bias = None info1.merge_batch(info2) self.assertEqual(info1.logit_bias.shape, (2, VOCAB_SIZE)) def test_merge_with_custom_logit_processor(self): """Test that merge combines processors when only one side has them.""" proc = MagicMock() info1 = _make_info(batch_size=1) info1.has_custom_logit_processor = True info1.custom_logit_processor = {1: (proc, torch.tensor([True]))} info1.custom_params = [{"a": 1}] info2 = _make_info(batch_size=1) info2.has_custom_logit_processor = False info2.custom_logit_processor = None info2.custom_params = None info1.merge_batch(info2) self.assertTrue(info1.has_custom_logit_processor) self.assertEqual(len(info1.custom_params), 2) def test_merge_with_none_sampling_seed(self): """Test that merge preserves None when both sampling_seeds are None.""" info1 = _make_info(batch_size=1) info1.sampling_seed = None info2 = _make_info(batch_size=1) info2.sampling_seed = None info1.merge_batch(info2) self.assertIsNone(info1.sampling_seed) def test_merge_with_both_sampling_seeds(self): """Test that merge concatenates both sampling_seed tensors.""" info1 = _make_info(batch_size=2) info1.sampling_seed = torch.tensor([10, 20], dtype=torch.int64) info2 = _make_info(batch_size=1) info2.sampling_seed = torch.tensor([30], dtype=torch.int64) info1.merge_batch(info2) self.assertEqual(info1.sampling_seed.shape[0], 3) self.assertEqual(info1.sampling_seed[0].item(), 10) self.assertEqual(info1.sampling_seed[1].item(), 20) self.assertEqual(info1.sampling_seed[2].item(), 30) # copy_for_forward class TestCopyForForward(CustomTestCase): def test_returns_copy_without_orchestrator(self): """Test that copy_for_forward returns a copy with orchestrator set to None.""" orch = MagicMock(is_required=False) info = _make_info(batch_size=1, penalizer_orchestrator=orch) copied = info.copy_for_forward() self.assertIsNone(copied.penalizer_orchestrator) # Original should still have orchestrator self.assertIsNotNone(info.penalizer_orchestrator) # from_schedule_batch class TestFromScheduleBatch(CustomTestCase): def _make_req( self, temp=1.0, top_p=1.0, top_k=-1, min_p=0.0, freq=0.0, presence=0.0, min_tokens=0, logit_bias=None, seed=None, stop_ids=None, eos_id=2, ): req = MagicMock() req.sampling_params.temperature = temp req.sampling_params.top_p = top_p req.sampling_params.top_k = top_k req.sampling_params.min_p = min_p req.sampling_params.frequency_penalty = freq req.sampling_params.presence_penalty = presence req.sampling_params.min_new_tokens = min_tokens req.sampling_params.logit_bias = logit_bias req.sampling_params.sampling_seed = seed req.sampling_params.stop_token_ids = stop_ids req.sampling_params.custom_params = None req.custom_logit_processor = None req.tokenizer.additional_stop_token_ids = None req.tokenizer.eos_token_id = eos_id return req @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_basic_construction(self, mock_server_args): """Test that from_schedule_batch correctly extracts sampling params from requests.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(temp=0.8, top_p=0.9, top_k=50, min_p=0.1)] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertEqual(len(info), 1) self.assertAlmostEqual(info.temperatures[0, 0].item(), 0.8, places=5) self.assertAlmostEqual(info.top_ps[0].item(), 0.9, places=5) self.assertEqual(info.top_ks[0].item(), 50) @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_mixed_batch_tensor_values_dtypes_shapes(self, mock_server_args): """All per-request params land in the right slot with the right dtype, and the fused float tensors are independent after filtering.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False params = [ dict(temp=0.5, top_p=0.7, top_k=10, min_p=0.2), dict(temp=1.0, top_p=1.0, top_k=1, min_p=0.0), dict(temp=1.3, top_p=0.95, top_k=TOP_K_ALL, min_p=0.05), ] batch = MagicMock() batch.reqs = [self._make_req(**p) for p in params] batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertEqual(info.temperatures.shape, (3, 1)) self.assertEqual(info.temperatures.dtype, torch.float32) self.assertEqual(info.top_ps.shape, (3,)) self.assertEqual(info.top_ps.dtype, torch.float32) self.assertEqual(info.min_ps.shape, (3,)) self.assertEqual(info.top_ks.dtype, torch.int32) for i, p in enumerate(params): self.assertAlmostEqual(info.temperatures[i, 0].item(), p["temp"], places=5) self.assertAlmostEqual(info.top_ps[i].item(), p["top_p"], places=5) self.assertEqual(info.top_ks[i].item(), p["top_k"]) self.assertAlmostEqual(info.min_ps[i].item(), p["min_p"], places=5) self.assertFalse(info.is_all_greedy) self.assertTrue(info.need_top_p_sampling) self.assertTrue(info.need_top_k_sampling) self.assertTrue(info.need_min_p_sampling) # Filtering must yield independent tensors with the right values even # though construction shares one fused buffer. keep = torch.tensor([2, 0], dtype=torch.int64) info.filter_batch([2, 0], keep) self.assertAlmostEqual(info.temperatures[0, 0].item(), 1.3, places=5) self.assertAlmostEqual(info.top_ps[1].item(), 0.7, places=5) self.assertEqual(info.top_ks[0].item(), TOP_K_ALL) @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_greedy_detection(self, mock_server_args): """Test that top_k=1 sets is_all_greedy=True.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(top_k=1)] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertTrue(info.is_all_greedy) @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_logit_bias_construction(self, mock_server_args): """Test that logit_bias dict is converted to a tensor with correct values.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(logit_bias={"5": 2.0, "10": -1.0})] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertIsNotNone(info.logit_bias) self.assertAlmostEqual(info.logit_bias[0, 5].item(), 2.0) self.assertAlmostEqual(info.logit_bias[0, 10].item(), -1.0) self.assertAlmostEqual(info.logit_bias[0, 0].item(), 0.0) @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_deterministic_seed(self, mock_server_args): """Test that explicit seed=123 is kept and missing seed defaults to 42.""" mock_server_args.return_value.enable_deterministic_inference = True mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(seed=123), self._make_req(seed=None)] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertIsNotNone(info.sampling_seed) self.assertEqual(info.sampling_seed[0].item(), 123) self.assertEqual(info.sampling_seed[1].item(), 42) # default @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_from_schedule_batch_sampling_flags(self, mock_server_args): """Test that sampling flags (need_top_p/top_k/min_p) are set correctly.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(top_p=0.9, top_k=50, min_p=0.1)] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertTrue(info.need_top_p_sampling) # 0.9 != 1.0 self.assertTrue(info.need_top_k_sampling) # 50 != TOP_K_ALL self.assertTrue(info.need_min_p_sampling) # 0.1 > 0 self.assertFalse(info.is_all_greedy) # top_k=50 > 1 @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_no_logit_bias_when_all_none(self, mock_server_args): """Test that logit_bias stays None when no request has logit_bias set.""" mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = False reqs = [self._make_req(), self._make_req()] batch = MagicMock() batch.reqs = reqs batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertIsNone(info.logit_bias) @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args") def test_custom_logit_processor_merging(self, mock_server_args): """Test deserialization and merging of custom logit processors.""" from sglang.srt.sampling.custom_logit_processor import ( DisallowedTokensLogitsProcessor, ) mock_server_args.return_value.enable_deterministic_inference = False mock_server_args.return_value.enable_custom_logit_processor = True proc_str = DisallowedTokensLogitsProcessor.to_str() req1 = self._make_req() req1.custom_logit_processor = proc_str req1.sampling_params.custom_params = {"token_ids": [1]} req2 = self._make_req() req2.custom_logit_processor = None # no processor req2.sampling_params.custom_params = None batch = MagicMock() batch.reqs = [req1, req2] batch.device = DEVICE info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE) self.assertTrue(info.has_custom_logit_processor) self.assertIsNotNone(info.custom_logit_processor) self.assertEqual(len(info.custom_logit_processor), 1) # Check the mask: req1 has processor (True), req2 doesn't (False) key = list(info.custom_logit_processor.keys())[0] proc, mask = info.custom_logit_processor[key] self.assertIsInstance(proc, DisallowedTokensLogitsProcessor) self.assertTrue(mask[0].item()) self.assertFalse(mask[1].item()) # custom_params should be collected for all reqs self.assertEqual(len(info.custom_params), 2) if __name__ == "__main__": unittest.main()