import torch

from losses.matching_loss import hungarian_matching_loss
from models.slot_extractor import SlotExtractor


def test_slot_extractor_and_matching_loss_backward():
    extractor = SlotExtractor(
        input_dim=32,
        state_dim=64,
        num_slots=4,
        num_layers=2,
        num_heads=4,
        dropout=0.0,
        num_subjects=12,
        num_predicates=7,
        num_objects=12,
    )
    tokens = torch.randn(2, 10, 32)
    mask = torch.ones(2, 10, dtype=torch.bool)
    predictions = extractor(tokens, mask)
    targets = [torch.tensor([[0, 0, 4], [1, 1, 5]]), torch.tensor([[1, 0, 6]])]
    result = hungarian_matching_loss(predictions, targets)
    assert torch.isfinite(result.loss)
    result.loss.backward()
    assert extractor.slot_queries.grad is not None