The 1st implementation of SGD works the second (which uses different backprop function)
does not, please fix the second.

Keep in mind:
- Dont tell me to use certain algorithm to initialize weights or normalize input, because
if 1st implenation does not require them second should not require it too.
- Learning rate is same in both case.
- I am using indirect way to compute gradient of softmax, once again this works 
fine in the 1st implementation, so it should not be a problem in the second too.
So, dont you dare tell me that I'm computing gradients wrong.

In short don't give me lame suggestions like regularize weights etc. because 1st 
implementation is able to converge nicely without anything fancy. There's some
issue that exists only in the second implementation, and they ony place where
1st and 2nd implementation differ is backpropagation and train function, spot that issue & fix.

1st implementation (works):
class DenseNeuralNetwork:
    def __init__(self, layer_sizes, activations, loss_f):
        self.learning_rate = None      # Learning rate
        self.layer_sizes = layer_sizes # Array storing the layer sizes
        self.loss_f = loss_f           # Loss function

        # The array of weights, biases, net inputs, and activation functions have their first element as None to represent that
        # input layers don't have them, this also helps match the indexes of expressions here match with their mathematical formulas

        # To store weights(W) and biases(b)
        np.random.seed(42)
        self.W = [None] + [np.random.uniform(-1, 1, (layer_sizes[i], layer_sizes[i-1])) for i in range(1, len(layer_sizes))]
        self.b = [None] + [np.random.uniform(-1, 1, layer_sizes[i]) for i in range(1, len(layer_sizes))]

        # To store activations(a) and net inputs(z)
        self.a = [np.zeros(size) for size in layer_sizes] # This one does not start with None, because input layer activates with same value as the inputs
        self.z = [None] + [np.zeros(layer_sizes[i]) for i in range(1, len(layer_sizes))]

        # Set the activation functions for each layer
        self.activation_f = activations # The first None element is passed as input here

    def propagate_forward(self, x):
        # Set input layer activations
        self.a[0] = x

        for i in range(1, len(self.W)):
            self.z[i] = self.W[i] @ self.a[i-1] + self.b[i]  # Weighted sum
            self.a[i] = self.activation_f[i].function(self.z[i])  # Activation

    def propagate_backward(self, y_true):
        error_term = {}

        # Store the index of the last layer as L
        L = len(self.W) - 1

        # For output layer:-
        if self.activation_f[L] == Softmax:
            error_term[L] = self.activation_f[L].derivative(self.z[L]) @ self.loss_f.gradient(y_true, self.a[L])
            # print('LAST LAYER ERROR TERM (FROM MATRIX MULTIPLICATION): ', error_term[L]) # For debugging
        else:
            error_term[L] = self.activation_f[L].derivative(self.z[L]) * self.loss_f.gradient(y_true, self.a[L])

        # Update weights and biases for output layer  based on those error terms
        self.W[L] -= self.learning_rate * np.outer(error_term[L], self.a[L-1])
        self.b[L] -= self.learning_rate * error_term[L]

        # For hidden layers:-
        for i in reversed(range(1, L)):
            ## Propgating error terms backward
            if  self.activation_f[i] == Softmax:
                error_term[i] = self.activation_f[i].derivative(self.z[i]) @ (self.W[i+1].T @ error_term[i+1])
            else:
                error_term[i] = self.activation_f[i].derivative(self.z[i]) * (self.W[i+1].T @ error_term[i+1])

            ## Updating the weights and biases based on those error terms
            self.W[i] -= self.learning_rate * np.outer(error_term[i], self.a[i-1])
            self.b[i] -= self.learning_rate * error_term[i]

    # Training loop
    def train(self, X_train, y_train, X_val, y_val, epochs=100, learning_rate=0.01, verbose=True):
        training_losses, validation_losses, training_accuracies, validation_accuracies = [], [], [], []
        self.learning_rate = learning_rate

        for epoch in range(epochs):
            total_loss, correct_predictions = 0, 0

            for i in range(X_train.shape[0]):
                x, y_true = X_train[i], y_train[i]
                self.propagate_forward(x)
                total_loss += self.loss_f.function(y_true, self.a[-1])

                if self.is_correct_prediction(self.a[-1], y_true):
                    correct_predictions += 1

                self.propagate_backward(y_true)

            avg_loss = total_loss / X_train.shape[0]
            accuracy = correct_predictions / X_train.shape[0] * 100
            training_losses.append(avg_loss)
            training_accuracies.append(accuracy)

            # Validation phase
            val_loss, val_correct_predictions = 0, 0
            for i in range(X_val.shape[0]):
                x, y_true = X_val[i], y_val[i]
                self.propagate_forward(x)
                val_loss += self.loss_f.function(y_true, self.a[-1])

                # Use np.argmax for validation comparison as well
                if self.is_correct_prediction(self.a[-1], y_true):
                    val_correct_predictions += 1

            val_avg_loss = val_loss / X_val.shape[0]
            val_accuracy = val_correct_predictions / X_val.shape[0] * 100
            validation_losses.append(val_avg_loss)
            validation_accuracies.append(val_accuracy)

            if verbose:
                print(f"Epoch {epoch+1}")
                print(f"Training Loss = {avg_loss}, Validation Loss = {val_avg_loss}")
                print(f"Training Accuracy: {accuracy:.3f}%, Validation Accuracy: {val_accuracy:.3f}%")
                print("----------------------------------------------------------------------")

        return training_losses, validation_losses, training_accuracies, validation_accuracies

    # Checks if the prediction has been made correctly; used to evalutue accuracy
    # This function may need to be altered based on the type of problem being solved
    def is_correct_prediction(self, predicted, true):
        if len(predicted) == 1:
            # For binary classification, if neuron o/p > 0.5 we consider it postive class
            return ((predicted > 0.5).astype(int)) == true
        else:
            # For multiclass classification
            # Use np.argmax to convert predicted and true values to class labels for comparison
            return np.argmax(predicted) == np.argmax(true)

    # For making predctions after training the network
    def predict(self, x):
        # Forward propagate through all layers
        self.propagate_forward(x)

        # The output of the last layer is the prediction
        y_pred = self.a[-1]

        return y_pred

    # For evaluating test accuracy later on
    def get_test_accuracy(self, X_test, y_test):
        correct_predictions = 0
        total_loss = 0

        for i in range(X_test.shape[0]):
            x, y_true = X_test[i], y_test[i]
            self.propagate_forward(x)

            # Accumulate loss
            total_loss += self.loss_f.function(y_true, self.a[-1])

            # Check if predicted class matches the true class
            if self.is_correct_prediction(self.a[-1], y_true):
                correct_predictions += 1

        # Calculate accuracy
        accuracy = correct_predictions / X_test.shape[0] * 100

        # Average loss
        avg_loss = total_loss / X_test.shape[0]

        print(f"Test Loss: {avg_loss}")
        print(f"Test Accuracy: {accuracy:.3f}%")

        return avg_loss, accuracy


2nd implementation (does not work):
import numpy as np

class DenseNeuralNetwork:
    def __init__(self, layer_sizes, activations, loss_f):
        self.learning_rate = None
        self.layer_sizes = layer_sizes
        self.loss_f = loss_f

        np.random.seed(42)
        self.W = [None] + [np.random.uniform(-1, 1, (layer_sizes[i], layer_sizes[i-1])) 
                           for i in range(1, len(layer_sizes))]
        self.b = [None] + [np.random.uniform(-1, 1, layer_sizes[i]) 
                           for i in range(1, len(layer_sizes))]

        self.a = [np.zeros(size) for size in layer_sizes]  # Activations
        self.z = [None] + [np.zeros(layer_sizes[i]) for i in range(1, len(layer_sizes))]  # Net inputs

        self.activation_f = activations  # Activation functions for each layer

    def propagate_forward(self, x):
        self.a[0] = x  # Input layer activation

        for i in range(1, len(self.W)):
            self.z[i] = self.W[i] @ self.a[i-1] + self.b[i]  # Weighted sum
            self.a[i] = self.activation_f[i].function(self.z[i])  # Activation

    def propagate_backward(self, y_true):
        error_term = {}
        dW = [None] + [np.zeros_like(w) for w in self.W[1:]]  # Gradient accumulators for weights
        db = [None] + [np.zeros_like(b) for b in self.b[1:]]  # Gradient accumulators for biases

        L = len(self.W) - 1  # Index of the last layer

        # Compute error term for the output layer
        if self.activation_f[L] == Softmax:
            # For Softmax with Categorical Cross-Entropy loss, the derivative simplifies to (a - y)
            error_term[L] = self.loss_f.gradient(y_true, self.a[L])
        else:
            error_term[L] = self.activation_f[L].derivative(self.z[L]) * self.loss_f.gradient(y_true, self.a[L])

        # Compute gradients for the output layer
        dW[L] = np.outer(error_term[L], self.a[L-1])  # Shape: (output_size, prev_layer_size)
        db[L] = error_term[L]  # Shape: (output_size,)

        # Compute error and gradients for hidden layers
        for i in reversed(range(1, L)):
            if self.activation_f[i] == Softmax:
                # Generally, Softmax is not used in hidden layers
                print(1)
                error = self.activation_f[i].derivative(self.z[i]) @ (self.W[i+1].T @ error_term[i+1])
            else:
                # For ReLU and others
                error = self.activation_f[i].derivative(self.z[i]) * (self.W[i+1].T @ error_term[i+1])
            
            error_term[i] = error  # Shape: (layer_size,)

            # Compute gradients for layer i
            dW[i] = np.outer(error_term[i], self.a[i-1])  # Shape: (current_layer_size, prev_layer_size)
            db[i] = error_term[i]  # Shape: (current_layer_size,)

        return dW, db  # Return gradients

    def train(self, X_train, y_train, X_val, y_val, epochs=100, learning_rate=0.01, verbose=True):
        training_losses, validation_losses, training_accuracies, validation_accuracies = [], [], [], []
        self.learning_rate = learning_rate

        n_samples = X_train.shape[0]
        for epoch in range(epochs):
            # Shuffle the training data at the beginning of each epoch
            permutation = np.random.permutation(n_samples)
            X_shuffled = X_train[permutation]
            y_shuffled = y_train[permutation]

            total_loss, correct_predictions = 0, 0

            for x, y_true in zip(X_shuffled, y_shuffled):
                self.propagate_forward(x)
                loss = self.loss_f.function(y_true, self.a[-1])
                total_loss += loss

                if self.is_correct_prediction(self.a[-1], y_true):
                    correct_predictions += 1

                # Get gradients
                dW, db = self.propagate_backward(y_true)

                # Update weights and biases immediately
                for layer in range(1, len(self.W)):
                    self.W[layer] -= self.learning_rate * dW[layer]
                    self.b[layer] -= self.learning_rate * db[layer]

            # Calculate average loss and accuracy for the epoch
            avg_loss = total_loss / n_samples
            accuracy = (correct_predictions / n_samples) * 100
            training_losses.append(avg_loss)
            training_accuracies.append(accuracy)

            # Validation phase
            val_loss, val_correct_predictions = 0, 0
            for i in range(X_val.shape[0]):
                x, y_true = X_val[i], y_val[i]
                self.propagate_forward(x)
                val_loss += self.loss_f.function(y_true, self.a[-1])

                if self.is_correct_prediction(self.a[-1], y_true):
                    val_correct_predictions += 1

            val_avg_loss = val_loss / X_val.shape[0]
            val_accuracy = (val_correct_predictions / X_val.shape[0]) * 100
            validation_losses.append(val_avg_loss)
            validation_accuracies.append(val_accuracy)

            if verbose:
                print(f"Epoch {epoch+1}")
                print(f"Training Loss = {avg_loss:.4f}, Validation Loss = {val_avg_loss:.4f}")
                print(f"Training Accuracy: {accuracy:.2f}%, Validation Accuracy: {val_accuracy:.2f}%")
                print("----------------------------------------------------------------------")

        return training_losses, validation_losses, training_accuracies, validation_accuracies

    def is_correct_prediction(self, predicted, true):
        if len(predicted) == 1:
            # Binary classification
            return ((predicted > 0.5).astype(int) == true).astype(int).sum()
        else:
            # Multiclass classification
            return np.argmax(predicted) == np.argmax(true)

    def predict(self, x):
        self.propagate_forward(x)
        y_pred = self.a[-1]
        return y_pred

    def get_test_accuracy(self, X_test, y_test):
        correct_predictions = 0
        total_loss = 0

        for i in range(X_test.shape[0]):
            x, y_true = X_test[i], y_test[i]
            self.propagate_forward(x)
            total_loss += self.loss_f.function(y_true, self.a[-1])

            if self.is_correct_prediction(self.a[-1], y_true):
                correct_predictions += 1

        accuracy = (correct_predictions / X_test.shape[0]) * 100
        avg_loss = total_loss / X_test.shape[0]

        print(f"Test Loss: {avg_loss:.4f}")
        print(f"Test Accuracy: {accuracy:.2f}%")

        return avg_loss, accuracy


Error with second implementation:
C:\Users\CircularCircle\AppData\Local\Temp\ipykernel_17532\3347165909.py:24: RuntimeWarning: overflow encountered in matmul
  self.z[i] = self.W[i] @ self.a[i-1] + self.b[i]  # Weighted sum
C:\Users\CircularCircle\AppData\Local\Temp\ipykernel_17532\1024956232.py:8: RuntimeWarning: invalid value encountered in subtract
  exp_z = np.exp(z - np.max(z))
Epoch 1
Training Loss = nan, Validation Loss = nan
Training Accuracy: 21.17%, Validation Accuracy: 24.50%
----------------------------------------------------------------------
Epoch 2
Training Loss = nan, Validation Loss = nan
Training Accuracy: 21.00%, Validation Accuracy: 24.50%
----------------------------------------------------------------------
Epoch 3
Training Loss = nan, Validation Loss = nan
Training Accuracy: 21.00%, Validation Accuracy: 24.50%