**FFNN (Feed Forward Neural Network)** - GeLU / ReLU / SiLU / SwiGLU

gistfile1.txt

import torch
import torch.nn as nn

# --------------------------------------------------
# Simple FFN that accepts a sentence
# --------------------------------------------------

class SimpleFFN(nn.Module):
    def __init__(self, activation="gelu"):
        super().__init__()

        # Tiny word embedding: each word -> 4 numbers
        self.embedding = nn.Embedding(100, 4)

        # FFN
        self.linear1 = nn.Linear(4, 8)
        self.linear2 = nn.Linear(8, 4)

        self.activation = activation

        # SwiGLU needs two branches
        if activation == "swiglu":
            self.gate = nn.Linear(4, 8)
            self.linear2 = nn.Linear(8, 4)

    def forward(self, x):

        # Word IDs -> vectors
        x = self.embedding(x)

        if self.activation == "relu":
            x = self.linear1(x)
            x = torch.relu(x)

        elif self.activation == "gelu":
            x = self.linear1(x)
            x = torch.nn.functional.gelu(x)

        elif self.activation == "silu":
            x = self.linear1(x)
            x = torch.nn.functional.silu(x)

        elif self.activation == "swiglu":
            # Main branch
            value = self.linear1(x)

            # Gate branch
            gate = self.gate(x)

            # SiLU + gating
            x = torch.nn.functional.silu(gate) * value

        x = self.linear2(x)

        return x


# --------------------------------------------------
# Sentence
# --------------------------------------------------

sentence = "the cat is sleeping"

# Give every word a simple ID
vocab = {
    "the": 0,
    "cat": 1,
    "is": 2,
    "sleeping": 3
}

word_ids = torch.tensor([[vocab[word] for word in sentence.split()]])

print("Words:", sentence.split())
print("Word IDs:", word_ids)


# --------------------------------------------------
# Try different FFN versions
# --------------------------------------------------

for activation in ["relu", "gelu", "silu", "swiglu"]:

    model = SimpleFFN(activation)

    output = model(word_ids)

    print("\nActivation:", activation)
    print("Output shape:", output.shape)
    print("Output:")
    print(output)