def test_sampling_determinism(model, device):
    """
    Some tests for sampling determinism.

    Note: DOES NOT test determinism for nucleus sampling, which is not currently
    compatible with "use_deterministic_algorithms" mode.
    """

    torch.use_deterministic_algorithms(True)

    def reset_seed():
        torch.manual_seed(0)
        np.random.seed(0)

    # ============================================================================
    # Argmax sampling (temperature=0) and categorical sampling with temperature=1
    # ============================================================================

    for temperature in [0, 1]:
        reset_seed()
        xf = get_test_model(model, device)

        # sample 1
        reset_seed()
        s1 = xf.sample("\n", num_tokens=10, temperature=temperature, top_p=None)

        # sample 2
        reset_seed()
        s2 = xf.sample("\n", num_tokens=10, temperature=temperature, top_p=None)

        assert (
            s1["strings"] == s2["strings"]
        ), f"Sampling for {model} was not deterministic with {temperature=}"

    Some tests for sampling determinism.

    Note: DOES NOT test determinism for nucleus sampling, which is not currently
    compatible with "use_deterministic_algorithms" mode.
    
neuron_explainer/tests/test_transformer.py