def test_masking_stats(tol: float = 1e-3):
    """Test to check that the masking probabilities
    match what we expect them to be.
    """
    kwargs = {
        "mask_prob": 0.15,
        "random_replace_prob": 0.1,
        "unmask_replace_prob": 0.1,
        "batch_size": 8,
    }
    tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    dataloader = create_data_iterator(**kwargs)
    num_samples = 10000
    total_tokens = 0
    masked_tokens = 0
    random_replace_tokens = 0
    unmasked_replace_tokens = 0
    for ix, batch in tqdm.tqdm(enumerate(dataloader, start=1), total=num_samples):
        # Since we don't mask the BOS / EOS tokens, we subtract them from the total tokens
        total_tokens += batch["attention_mask"].sum().item() - (
            2 * batch["attention_mask"].size(0)
        )
        masked_tokens += (batch["tgt_tokens"] != tokenizer.pad_token_id).sum().item()
        random_or_unmasked = (
            batch["tgt_tokens"] != tokenizer.pad_token_id
        ).logical_and(batch["src_tokens"] != tokenizer.mask_token_id)
        unmasked = random_or_unmasked.logical_and(
            batch["src_tokens"] == batch["tgt_tokens"]
        )
        unmasked_replace_tokens += unmasked.sum().item()
        random_replace_tokens += random_or_unmasked.sum().item() - unmasked.sum().item()
        if ix == num_samples:
            break
    estimated_mask_prob = masked_tokens / total_tokens
    estimated_random_tokens = random_replace_tokens / total_tokens
    estimated_unmasked_tokens = unmasked_replace_tokens / total_tokens
    assert np.isclose(estimated_mask_prob, kwargs["mask_prob"], atol=tol)
    assert np.isclose(
        estimated_random_tokens,
        kwargs["random_replace_prob"] * kwargs["mask_prob"],
        atol=tol,
    )
    assert np.isclose(
        estimated_unmasked_tokens,
        kwargs["unmask_replace_prob"] * kwargs["mask_prob"],
        atol=tol,
    )
Test to check that the masking probabilities
    match what we expect them to be.
    
training/HelloDeepSpeed/tests/test_train_bert.py