Code:
# Softmax activation function
class Softmax:
    @staticmethod
    def function(z):
        # To avoid overflow issues with large values of z, we can subtract the maximum value
        # from the input vector z. This doesn’t affect the result because the same constant is
        # subtracted from all elements, and the softmax formula is invariant to this shift
        exp_z = np.exp(z - np.max(z))
        return exp_z / np.sum(exp_z)

    @staticmethod
    def derivative(z):
        # Unlike other activation functions, the result of Softmax depends on all neurons in the layer,
        # so during backprop we take gradient of it's results(a) w.r.t. each neuron's net-inputs(z) in that layer,
        # doing that for each activation results in a matrix (instead of a vector, for other activations)
        softmax_val = Softmax.function(z)
        jacobian = np.diag(softmax_val) - np.outer(softmax_val, softmax_val)
        return jacobian

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:
                #print(2)
                # 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

        for epoch in range(epochs):
            # Initialize accumulators for gradients
            dW_accum = [None] + [np.zeros_like(w) for w in self.W[1:]]
            db_accum = [None] + [np.zeros_like(b) for b in self.b[1:]]
            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

                # Get gradients instead of updating weights
                dW, db = self.propagate_backward(y_true)

                # Accumulate gradients
                for layer in range(1, len(self.W)):
                    dW_accum[layer] += dW[layer]
                    db_accum[layer] += db[layer]

            # Average loss and accuracy
            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)

            # Update weights and biases using averaged gradients
            for layer in range(1, len(self.W)):
                self.W[layer] -= self.learning_rate * (dW_accum[layer] / X_train.shape[0])
                self.b[layer] -= self.learning_rate * (db_accum[layer] / X_train.shape[0])

            # 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:
Epoch 1
Training Loss = 8.9916, Validation Loss = 20.8614
Training Accuracy: 54.00%, Validation Accuracy: 24.50%
----------------------------------------------------------------------
Epoch 2
Training Loss = 21.8285, Validation Loss = 12.0195
Training Accuracy: 21.00%, Validation Accuracy: 56.50%
----------------------------------------------------------------------
Epoch 3
Training Loss = 11.8353, Validation Loss = nan
Training Accuracy: 57.17%, Validation Accuracy: 24.50%
----------------------------------------------------------------------
Epoch 4
Training Loss = nan, Validation Loss = nan
Training Accuracy: 21.00%, Validation Accuracy: 24.50%
----------------------------------------------------------------------
C:\Users\CircularCircle\AppData\Local\Temp\ipykernel_13896\2064518748.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_13896\1024956232.py:8: RuntimeWarning: invalid value encountered in subtract
  exp_z = np.exp(z - np.max(z))