Gradient Descent Beginner

import numpy as np

# Full implementation of mini-batch gradient descent
def mini_batch_gd(X, y, lr=0.01, batch_size=32, epochs=100):
    n_samples, n_features = X.shape
    w = np.zeros(n_features)
    b = 0.0
    losses = []

    for epoch in range(epochs):
        # Shuffle data each epoch
        indices = np.random.permutation(n_samples)
        X_shuffled = X[indices]
        y_shuffled = y[indices]

        for i in range(0, n_samples, batch_size):
            X_batch = X_shuffled[i:i+batch_size]
            y_batch = y_shuffled[i:i+batch_size]

            # Forward pass
            predictions = X_batch @ w + b
            error = predictions - y_batch

            # Compute gradients
            dw = (2 / len(X_batch)) * X_batch.T @ error
            db = (2 / len(X_batch)) * np.sum(error)

            # Update parameters
            w -= lr * dw
            b -= lr * db

        loss = np.mean((X @ w + b - y) ** 2)
        losses.append(loss)

    return w, b, losses

def sgd_momentum(gradient_fn, x0, lr=0.01, momentum=0.9, n_steps=100):
    x = x0.copy()
    velocity = np.zeros_like(x)

    for _ in range(n_steps):
        grad = gradient_fn(x)
        velocity = momentum * velocity - lr * grad  # Accumulate
        x = x + velocity                            # Update

    return x