def test_mask_random_tokens(self):
    """Test only standard tokens.

    This test deals with sentences composed of two parts: [100 CLS tokens, 100
    chosen standard tokens]. CLS is the token that is added at the beginning of
    the sentence and there is only one token in standard scenario. It is never
    masked because it is not a part of the sentence.
    This tests whether mask_random_tokens will:
      - mask only standard tokens
      - mask expected number of tokens (15 percent candidates for masking)
    """
    cls_token = 101
    mask_token = 103
    example_standard_token = 1001
    test_case_row = np.array([cls_token] * 100 + [example_standard_token] * 100)
    test_case = [(test_case_row.copy(),)]

    out, original_tokens, token_weights = next(
        tf_inputs.mask_random_tokens(test_case))
    # test whether original tokens are unchanged
    self.assertAllEqual(test_case_row, original_tokens)

    self.assertEqual(1, token_weights.sum())
    self.assertEqual(
        15,
        (token_weights > 0).sum())  # we should have 15 candidates for masking

    # 101 is a special token, so only 1001 should be masked
    self.assertAllEqual(out[:100], test_case_row[:100])

    # Each candidate has 0.8 probability to be masked while others have 0, so
    # no more than 15 tokens with MASK
    self.assertLessEqual((out == mask_token).sum(), 15)
Test only standard tokens.

    This test deals with sentences composed of two parts: [100 CLS tokens, 100
    chosen standard tokens]. CLS is the token that is added at the beginning of
    the sentence and there is only one token in standard scenario. It is never
    masked because it is not a part of the sentence.
    This tests whether mask_random_tokens will:
      - mask only standard tokens
      - mask expected number of tokens (15 percent candidates for masking)
    
trax/data/tf_inputs_test.py