# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/00_core.ipynb (unless otherwise specified).

__all__ = ['max_grad', 'semiring', 'Log', 'Max', 'grad', 'SequenceDist']

# Cell
import torch
from collections import namedtuple

# Cell
semiring = namedtuple('semiring', ('zero', 'one', 'mul', 'sum', 'dsum'))

def max_grad(x, dim=0):
    return torch.zeros_like(x).scatter_(dim, x.argmax(dim, True), 1.0)

Log = semiring(zero=-1e38, one=0., mul=torch.add, sum=torch.logsumexp, dsum=torch.softmax)
Max = semiring(zero=-1e38, one=0., mul=torch.add, sum=(lambda x, dim=0: torch.max(x, dim=dim)[0]), dsum=max_grad)

# Cell
def grad(f, x):
    x = x.detach().requires_grad_()
    with torch.enable_grad():
        y = f(x)
    return torch.autograd.grad(y, x)[0].detach()

class SequenceDist():
    def __init__(self):
        pass

    def logZ(self, scores, S:semiring=Log):
        raise NotImplementedError

    def viterbi(self, scores):
        raise NotImplementedError

    def ctc_loss(self, scores, targets, target_lengths):
        raise NotImplementedError

    def posteriors(self, scores, S:semiring=Log):
        f = lambda x: self.logZ(x, S).sum()
        return grad(f, scores)