Bengio layers

layers.py

class Embedding():
    def __init__(self, device, num_embeddings, embedding_dim):
        self.out = None
        self.weight = torch.randn(num_embeddings, embedding_dim, device=device)

    def __call__(self, x):
        self.out = F.embedding(x, self.weight)
        return self.out

    def params(self):
        return [self.weight]

class Flatten():
    def __init__(self, input_dim1, input_dim2):
        self.out = None
        self.input_dim1 = input_dim1
        self.input_dim2 = input_dim2

    def __call__(self, x):
        assert x.ndim == 3
        self.out = x.view(-1, self.input_dim1 * self.input_dim2)
        return self.out

    def params(self):
        return []

class Linear():
    def __init__(self, device, in_features, out_features, bias=True):
        self.out = None
        self.weight = torch.randn(in_features, out_features, device=device)
        self.bias = torch.zeros(out_features, device=device) if bias else None

    def __call__(self, x):
        self.out = x @ self.weight
        if self.bias is not None:
            self.out += self.bias

        return self.out

    def params(self):
        return [self.weight] + ([] if self.bias is None else [self.bias])

class Tanh():
    def __init__(self):
        self.out = None

    def __call__(self, x):
        self.out = torch.tanh(x)
        return self.out

    def params(self):
        return []