import torch

from verl.utils.torch_functional import logprobs_from_logits_v2


def test_hint_next_token_labels_match_causal_logits():
    batch_size = 3
    sequence_length = 8
    vocab_size = 16
    input_ids = torch.randint(vocab_size, (batch_size, sequence_length))
    responses = input_ids[:, 1:]
    logits = torch.randn(batch_size, sequence_length, vocab_size)
    causal_logits = logits[:, -responses.size(1) - 1:-1, :]

    assert causal_logits.shape[:-1] == responses.shape
    assert logprobs_from_logits_v2(causal_logits, responses).shape == responses.shape