# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/old/00_core.ipynb.

# %% auto 0
__all__ = ['Embedding', 'sigmoid', 'SigmoidFunction', 'tanh', 'TanhFunction', 'relu', 'ReLUFunction', 'CrossEntropyLoss']

# %% ../../nbs/old/00_core.ipynb 10
class Embedding(LayerWithCache):
    def __init__(self, vocab_size, emb_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.emb_size = emb_size
        self.E = np.random.normal(loc=0, scale=0.1, size=(vocab_size, emb_size))
        self.params = (self.E, )


    def forward(self, x):
        out = self.E[x, :]
        self.to_cache(x)
        return out

    def backward(self, grad):
        x = self.from_cache()
        dE = np.zeros_like(self.E)
        np.add.at(dE, x, grad)
        return dE

    def __call__(self, x):
        return self.forward(x)

    def parameters(self):
        return self.params

    def update(self, params):
        self.E = params[0]

    def reset_cache(self):
        self.reset()


# %% ../../nbs/old/00_core.ipynb 12
def sigmoid(x):
    s = 1.0 / (1.0 + np.exp(-x))
    return s

# %% ../../nbs/old/00_core.ipynb 13
class SigmoidFunction(LayerWithCache):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        a = sigmoid(x)
        self.to_cache(a)
        return a

    def backward(self, grad):
        a = self.from_cache()
        return a * (1. - a) * grad.reshape(a.shape)

    def __call__(self, x):
        return self.forward(x)

    def reset_cache(self):
        self.reset()


# %% ../../nbs/old/00_core.ipynb 14
def tanh(x):
    return (np.exp(x) - np.exp(-x)) / (np.exp(x) + np.exp(-x))

# %% ../../nbs/old/00_core.ipynb 15
class TanhFunction(LayerWithCache):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        a = np.tanh(x)
        self.to_cache(a)
        return a

    def backward(self, grad):
        a = self.from_cache()
        return (1. - a ** 2) * grad.reshape(a.shape)

    def __call__(self, x):
        return self.forward(x)

    def reset_cache(self):
        self.reset()


# %% ../../nbs/old/00_core.ipynb 16
def relu(x):
    return np.maximum(x, 0)

# %% ../../nbs/old/00_core.ipynb 17
class ReLUFunction(LayerWithCache):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        a = relu(x)
        self.to_cache(a)
        return a

    def backward(self, grad):
        dx = copy(grad)
        a = self.from_cache()
        dx[a <= 0] = 0
        return dx

    def __call__(self, x):
        return self.forward(x)

# %% ../../nbs/old/00_core.ipynb 19
class CrossEntropyLoss(LayerWithCache):
    def __init__(self, eps=1e-8):
        super().__init__()
        self.eps = eps


    def forward(self, output, target):
        num_classes = output.shape[-1]
        target = one_hot_encoder(target, num_classes)
        target = target.reshape(-1, num_classes)
        self.to_cache((output, target))
        logs = np.log(output + self.eps)
        loss = np.multiply(-target, logs).sum(axis=-1)
        return loss.mean()

    def backward(self):
        output, target = self.from_cache()
        return output - target

    def __call__(self, output, target):
        return self.forward(output, target)

    def reset_cache(self):
        self.reset()

