def test_key_values_correct(
      self,
      test_dict,
      mask_keys,
      expected_dict,
  ):
    """Match masked dict with expected value."""
    masked_dict = optim_utils.create_dict_mask(test_dict, mask_keys)
    self.assertDictEqual(masked_dict, expected_dict)
Match masked dict with expected value.
language/mentionmemory/utils/optim_utils_test.py