# Source: content/notes/deep-learning/rnns-and-sequence-models.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import math
import torch
from torch import nn

torch.manual_seed(8)
torch.set_num_threads(1)
logits = torch.zeros(3, 1, 2, requires_grad=True)
loss_fn = nn.CTCLoss(blank=0, reduction="sum", zero_infinity=False)
loss = loss_fn(logits.log_softmax(-1), torch.tensor([1]),
               torch.tensor([3]), torch.tensor([1]))
torch.testing.assert_close(loss, torch.tensor(-math.log(0.75)))
loss.backward()
assert logits.grad is not None and torch.isfinite(logits.grad).all()
bad = loss_fn(torch.zeros(2, 1, 2).log_softmax(-1), torch.tensor([1, 1]),
              torch.tensor([2]), torch.tensor([2]))
assert torch.isinf(bad)
good = loss_fn(torch.zeros(3, 1, 2).log_softmax(-1), torch.tensor([1, 1]),
               torch.tensor([3]), torch.tensor([2]))
torch.testing.assert_close(good, torch.tensor(-math.log(0.125)))
print({"one_a_nll": loss.item(), "two_a_nll": good.item()})