Neural Networks

import numpy as np

def relu(z):
    return np.maximum(0, z)

def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)   # numerical stability
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

class MLP:
    def __init__(self, in_dim, hidden_dim, out_dim, seed=0):
        rng = np.random.default_rng(seed)
        self.W1 = rng.normal(0, 0.1, (hidden_dim, in_dim))
        self.b1 = np.zeros(hidden_dim)
        self.W2 = rng.normal(0, 0.1, (out_dim, hidden_dim))
        self.b2 = np.zeros(out_dim)

    def forward(self, x):                # x: (batch, in_dim)
        z1 = x @ self.W1.T + self.b1     # (batch, hidden_dim)
        a1 = relu(z1)
        z2 = a1 @ self.W2.T + self.b2    # (batch, out_dim)
        return softmax(z2)

mlp = MLP(in_dim=784, hidden_dim=128, out_dim=10)
x = np.random.randn(4, 784)
probs = mlp.forward(x)
print(probs.shape, probs.sum(axis=-1))   # (4, 10) [1. 1. 1. 1.]