def test_sample_utilities():
    """Correctness tests for functions that enable batched sampling."""

    # ========================
    # test prep_input_and_pad
    # ========================
    # make a set of test prompts of unequal lengths
    test_inputs = [
        [1, 100, 50, 47],  # prompt 0
        [1298, 618, 952, 223, 4, 42],  # prompt 1
        [31],  # prompt 2
    ]
    X_ref = torch.LongTensor(
        [
            [0, 0, 1, 100, 50, 47],
            [1298, 618, 952, 223, 4, 42],
            [0, 0, 0, 0, 0, 31],
        ]
    )
    pad_ref = torch.BoolTensor(
        [
            [1, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0],
            [1, 1, 1, 1, 1, 0],
        ]
    )
    X, pad = prep_input_and_pad(test_inputs, pad_side="left", device="cpu")

    assert_equal(X_ref, X, msg="prep_input_and_pad produced incorrect X")
    assert_equal(pad_ref, pad, msg="prep_input_and_pad produced incorrect pad")

    # =====================================
    # test prep_pos_from_pad_and_prev_lens
    # =====================================
    # based on the same example as before, compute the associated pos

    pos_ref = torch.LongTensor([[0, 0, 0, 1, 2, 3], [0, 1, 2, 3, 4, 5], [0, 0, 0, 0, 0, 0]])
    prev_lens = torch.zeros(3, 1).long()
    pos = prep_pos_from_pad_and_prev_lens(pad_ref, prev_lens)

    assert_equal(pos_ref, pos, msg="prep_pos_from_pad_and_prev_lens produced incorrect pos")

    # simulate one step forward of sampling
    seq_lens = (pos[:, -1] + 1).unsqueeze(-1)
    new_pad_ref = torch.BoolTensor(
        [
            [0],
            [0],
            [0],
        ]
    )
    new_pos = prep_pos_from_pad_and_prev_lens(new_pad_ref, seq_lens)

    new_pos_ref = torch.LongTensor([[4], [6], [1]])
    assert_equal(new_pos_ref, new_pos, msg="prep_pos_from_pad_and_prev_lens produced incorrect pos")
Correctness tests for functions that enable batched sampling.
neuron_explainer/tests/test_transformer.py