import torch


def test_chg_call_matches_forward(llama_chg):
    input_ids = torch.tensor([[1, 2, 3]], device=llama_chg.device)

    out_forward = llama_chg.forward(input_ids)
    out_call = llama_chg(input_ids)

    assert torch.equal(out_forward, out_call)


def test_chg_masks_shapes_and_effect(llama_chg):
    input_ids = torch.tensor([[1, 2]], device=llama_chg.device)

    ones_mask = torch.ones(llama_chg.num_layers, llama_chg.num_heads)
    zeros_mask = torch.zeros(llama_chg.num_layers, llama_chg.num_heads)

    out_ones = llama_chg(input_ids, ones_mask)
    out_zeros = llama_chg(input_ids, zeros_mask)

    assert out_ones.shape == out_zeros.shape
    assert out_ones.shape[:2] == (1, 2)
    assert not torch.allclose(out_ones, out_zeros)

    multi_masks = torch.stack([ones_mask, zeros_mask], dim=0)
    out_multi = llama_chg(input_ids, multi_masks)
    assert out_multi.shape[0] == 1
    assert out_multi.shape[1] == 2


def test_chg_create_masks_shapes(llama_chg):
    masks = llama_chg.create_masks(num_masks=None, as_parameter=False, as_logits=False)
    assert masks.shape == (llama_chg.num_layers, llama_chg.num_heads)
    assert torch.all((masks >= 0) & (masks <= 1))

    masks_multi = llama_chg.create_masks(num_masks=3, as_parameter=False, as_logits=False)
    assert masks_multi.shape == (3, llama_chg.num_layers, llama_chg.num_heads)
