diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 887c04e..9cee71c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -1,5 +1,4 @@
-# .github/workflows/ci.yml
-name: C/C++ CI
+name: Neural Network CI
 
 on:
   pull_request:
@@ -12,17 +11,29 @@ jobs:
     steps:
       - uses: actions/checkout@v2
 
+      - name: Cache build dependencies
+        uses: actions/cache@v2
+        with:
+          path: |
+            ~/.cache/pip
+            build
+          key: ${{ runner.os }}-build-${{ hashFiles('**/Makefile') }}
+          restore-keys: |
+            ${{ runner.os }}-build-
+
       - name: Install build tools and linter
-        run: sudo apt-get update && sudo apt-get install -y clang-format cppcheck
+        run: sudo apt-get update && sudo apt-get install -y clang-format cppcheck libcunit1-dev
 
       - name: Check code formatting
         run: |
-          mkdir -p ../temp_repo
-          cp -r . ../temp_repo
-          find ../temp_repo -name "*.c" -o -name "*.h" | xargs clang-format -i --style=Google
-          diff -u -r ../temp_repo .
+          find . -name "*.c" -o -name "*.h" | xargs clang-format --dry-run --Werror --style=Google
+
+      - name: Build project
+        run: make all
+
+      - name: Run tests
+        run: make test
 
       - name: Run static code analysis
         run: |
-          # The suppress flag tells cppcheck to ignore unused function warnings.
-          cppcheck --enable=all -I nn/include --suppress=missingIncludeSystem --suppress=unusedFunction --inconclusive --error-exitcode=1 .
\ No newline at end of file
+          cppcheck --enable=all -I nn/include --suppress=missingIncludeSystem --inconclusive --error-exitcode=1 .
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 6c1d726..eb0ba0b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,7 @@
 build
 
 # Doxygen output
-docs/api/
\ No newline at end of file
+docs/api/
+
+# Object files
+*.o
\ No newline at end of file
diff --git a/Makefile b/Makefile
new file mode 100644
index 0000000..c322603
--- /dev/null
+++ b/Makefile
@@ -0,0 +1,44 @@
+# Compiler and flags
+CC = gcc
+CFLAGS = -I nn/include -I tests -Wall -Wextra -Werror -Wpedantic -Wstrict-prototypes -Wold-style-definition -g $(CU_CFLAGS)
+
+# Source files
+SRCS = $(shell find nn/src -name '*.c' -not -path 'nn/src/main.c')
+TEST_SRCS = tests/core_cunit.c tests/nn_cunit.c tests/cunit_runner.c tests/test_utils.c
+
+# Object files
+OBJS = $(SRCS:.c=.o)
+TEST_OBJS = $(TEST_SRCS:.c=.o)
+
+# Executables
+EXAMPLE_SRCS = $(wildcard nn/src/examples/*.c)
+EXAMPLE_OBJS = $(EXAMPLE_SRCS:.c=.o)
+EXAMPLE_EXECS = $(EXAMPLE_SRCS:.c=)
+TEST_RUNNER = tests/cunit_runner_exec
+
+# CUnit flags
+CU_CFLAGS = $(shell pkg-config --cflags cunit)
+CU_LIBS = $(shell pkg-config --libs cunit)
+
+# Targets
+all: $(EXAMPLE_EXECS) $(TEST_RUNNER)
+
+$(EXAMPLE_EXECS): nn/src/examples/%: $(filter-out nn/src/examples/%.o, $(OBJS)) nn/src/examples/%.o
+	$(CC) $(CFLAGS) -o $@ $^ -lm
+
+$(TEST_RUNNER): $(filter-out $(EXAMPLE_OBJS), $(OBJS)) $(TEST_OBJS)
+	$(CC) $(CFLAGS) $(CU_CFLAGS) -o $@ $^ $(CU_LIBS) -lm
+
+%.o: %.c
+	$(CC) $(CFLAGS) -c -o $@ $<
+
+clean:
+	rm -f $(OBJS) $(TEST_OBJS) $(EXAMPLE_OBJS) $(EXAMPLE_EXECS) $(TEST_RUNNER)
+
+test: $(TEST_RUNNER)
+	./$(TEST_RUNNER)
+
+format:
+	clang-format -i $(SRCS) $(TEST_SRCS) nn/include/*.h
+
+.PHONY: all clean test format
\ No newline at end of file
diff --git a/nn/include/activation.h b/nn/include/activation.h
index 9306eef..3e6bdae 100644
--- a/nn/include/activation.h
+++ b/nn/include/activation.h
@@ -1,121 +1,177 @@
-#pragma once
-
-#include "linalg.h"
-
 /**
  * @file activation.h
- * @brief Activation functions and their derivatives operating on matrices.
+ * @brief Header for activation functions used in neural networks.
  *
- * All functions return newly allocated matrices; callers own and must free
- * the results. Derivatives are provided for use in backpropagation.
+ * This file defines an enumeration for various activation functions and
+ * declares their corresponding matrix-based functions and their derivatives.
+ * These functions are crucial for introducing non-linearity into neural
+ * networks.
  */
 
-//============================
-// Activation Functions
-//============================
+#ifndef NN_ACTIVATION_H
+#define NN_ACTIVATION_H
+
+#include <stddef.h>
+
+#include "linalg.h"  // Assumes Matrix struct is defined here
+
+/**
+ * @brief Enum for different activation functions.
+ *
+ * This enumeration lists the types of activation functions supported
+ * within the neural network framework. Each enumerator corresponds to
+ * a specific non-linear function applied to the output of a layer.
+ */
+typedef enum {
+  RELU,       /**< Rectified Linear Unit activation. */
+  SIGMOID,    /**< Sigmoid activation. */
+  SOFTMAX,    /**< Softmax activation, typically used in the output layer for
+                 multi-class classification. */
+  TANH,       /**< Hyperbolic Tangent activation. */
+  LEAKY_RELU, /**< Leaky Rectified Linear Unit activation. */
+  SIGN,       /**< Sign activation. */
+  IDENTITY,   /**< Identity activation. */
+  HARD_TANH   /**< Hard Tanh activation. */
+} activation_function;
+
+// Activation functions
 
 /**
- * @brief Sigmoid activation applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with sigmoid applied.
+ * @brief Applies the Sigmoid activation function element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the sigmoid function applied to each element.
  */
 Matrix* sigmoid(Matrix* m);
+
 /**
- * @brief Derivative of sigmoid applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with sigmoid derivative applied.
+ * @brief Applies the ReLU (Rectified Linear Unit) activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the ReLU function applied to each element.
  */
-Matrix* sigmoid_prime(Matrix* m);
+Matrix* relu(Matrix* m);
 
 /**
- * @brief ReLU activation applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with ReLU applied.
+ * @brief Applies the Hyperbolic Tangent (tanh) activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the tanh function applied to each element.
  */
-Matrix* relu(Matrix* m);
+Matrix* tanh_activation(Matrix* m);
+
 /**
- * @brief Derivative of ReLU applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with ReLU derivative applied.
+ * @brief Applies the Leaky ReLU activation function element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @param leak_parameter The leak parameter (alpha) for the Leaky ReLU.
+ * @return A new Matrix with the Leaky ReLU function applied to each element.
  */
-Matrix* relu_prime(Matrix* m);
+Matrix* leaky_relu(Matrix* m, double leak_parameter);
 
 /**
- * @brief Hyperbolic tangent activation applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with tanh applied.
+ * @brief Applies the Sign activation function element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the Sign function applied to each element.
  */
-Matrix* tanh_activation(Matrix* m);
+Matrix* sign_activation(Matrix* m);
+
 /**
- * @brief Derivative of tanh applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with tanh derivative applied.
+ * @brief Applies the Identity activation function element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix that is a copy of the input matrix.
  */
-Matrix* tanh_prime(Matrix* m);
+Matrix* identity_activation(Matrix* m);
 
 /**
- * @brief Leaky ReLU activation (alpha=0.01) applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with Leaky ReLU applied.
+ * @brief Applies the Hard Tanh activation function element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the Hard Tanh function applied to each element.
  */
-Matrix* leaky_relu(Matrix* m);
+Matrix* hard_tanh(Matrix* m);
+
 /**
- * @brief Derivative of Leaky ReLU (alpha=0.01) applied elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with Leaky ReLU derivative applied.
+ * @brief Applies the Softmax activation function to a matrix.
+ * This function is typically used in the output layer of a neural network for
+ * multi-class classification. It normalizes the input values into a probability
+ * distribution.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the Softmax function applied.
  */
-Matrix* leaky_relu_prime(Matrix* m);
+Matrix* softmax(Matrix* m);
+
+// Derivatives of activation functions
 
 /**
- * @brief Leaky ReLU with custom alpha (leak_parameter).
- * @param m Input matrix (m x n).
- * @param leak_parameter The alpha value for the leak.
- * @return New matrix (m x n) with Leaky ReLU applied.
+ * @brief Computes the derivative of the Sigmoid activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix (output of the sigmoid function).
+ * @return A new Matrix with the sigmoid derivative applied to each element.
  */
-Matrix* leaky_relu_with_alpha(Matrix* m, double leak_parameter);
+Matrix* sigmoid_prime(Matrix* m);
+
 /**
- * @brief Derivative of Leaky ReLU with custom alpha.
- * @param m Input matrix (m x n).
- * @param leak_parameter The alpha value for the leak.
- * @return New matrix (m x n) with Leaky ReLU derivative applied.
+ * @brief Computes the derivative of the ReLU activation function element-wise
+ * to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the ReLU derivative applied to each element.
  */
-Matrix* leaky_relu_prime_with_alpha(Matrix* m, double leak_parameter);
+Matrix* relu_prime(Matrix* m);
 
 /**
- * @brief Sign activation: -1, 0, or +1 elementwise.
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with sign activation applied.
+ * @brief Computes the derivative of the Hyperbolic Tangent (tanh) activation
+ * function element-wise to a matrix.
+ * @param m A pointer to the input Matrix (output of the tanh function).
+ * @return A new Matrix with the tanh derivative applied to each element.
  */
-Matrix* sign_activation(Matrix* m);
+Matrix* tanh_prime(Matrix* m);
+
 /**
- * @brief Derivative of sign (0 everywhere; undefined at 0).
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) of zeros.
+ * @brief Computes the derivative of the Leaky ReLU activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @param leak_parameter The leak parameter (alpha) for the Leaky ReLU.
+ * @return A new Matrix with the Leaky ReLU derivative applied to each element.
  */
-Matrix* sign_prime(Matrix* m);
+Matrix* leaky_relu_prime(Matrix* m, double leak_parameter);
 
 /**
- * @brief Identity activation (returns a copy).
- * @param m Input matrix (m x n).
- * @return New matrix (m x n), a copy of the input.
+ * @brief Computes the derivative of the Sign activation function element-wise
+ * to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the Sign derivative applied to each element.
  */
-Matrix* identity_activation(Matrix* m);
+Matrix* sign_prime(Matrix* m);
+
 /**
- * @brief Derivative of identity (ones).
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) of ones.
+ * @brief Computes the derivative of the Identity activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix of the same size as the input, with all elements set
+ * to 1.0.
  */
 Matrix* identity_prime(Matrix* m);
 
 /**
- * @brief Hard Tanh activation clamped to [-1, 1].
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with hard tanh applied.
+ * @brief Computes the derivative of the Hard Tanh activation function
+ * element-wise to a matrix.
+ * @param m A pointer to the input Matrix.
+ * @return A new Matrix with the Hard Tanh derivative applied to each element.
  */
-Matrix* hard_tanh(Matrix* m);
+Matrix* hard_tanh_prime(Matrix* m);
+
 /**
- * @brief Derivative of Hard Tanh (1 in (-1,1), else 0).
- * @param m Input matrix (m x n).
- * @return New matrix (m x n) with hard tanh derivative applied.
+ * @brief Computes the derivative of the Softmax activation function
+ * element-wise to a matrix. This is typically used in conjunction with a loss
+ * function like cross-entropy where the derivative simplifies to output * (1 -
+ * output).
+ * @param m A pointer to the input Matrix (output of the softmax function).
+ * @return A new Matrix with the Softmax derivative applied to each element.
  */
-Matrix* hard_tanh_prime(Matrix* m);
+Matrix* softmax_prime(Matrix* m);
+
+#endif  // NN_ACTIVATION_H
+/**
+ * @brief Converts an activation function enum to its string representation.
+ * @param func The activation function enum.
+ * @return A string representing the activation function.
+ */
+const char* activation_to_string(activation_function func);
\ No newline at end of file
diff --git a/nn/include/backprop.h b/nn/include/backprop.h
index 29aa941..b911023 100644
--- a/nn/include/backprop.h
+++ b/nn/include/backprop.h
@@ -16,11 +16,11 @@
  * @brief Compute layer-wise deltas and cache them for gradient evaluation.
  * @param nn Pointer to network.
  * @param y_true Ground-truth labels/targets.
- * @param loss_func Loss function (optional here; for monitoring).
+ * @param loss_type Loss function (optional here; for monitoring).
  * @param loss_func_grad Gradient of loss w.r.t predictions (required).
  */
 void backpropagate(NeuralNetwork* nn, const Matrix* y_true,
-                   LossFunction loss_func, LossFunctionGrad loss_func_grad);
+                   LossFunctionType loss_type, LossFunctionGrad loss_func_grad);
 
 /** @brief Calculate weight gradient for a specific layer. */
 Matrix* calculate_weight_gradient(const Cache* cache, size_t layer_index,
diff --git a/nn/include/cache.h b/nn/include/cache.h
index 1057f14..f0e7319 100644
--- a/nn/include/cache.h
+++ b/nn/include/cache.h
@@ -17,8 +17,11 @@ typedef struct Cache Cache;
 /** @brief Create an empty cache. */
 Cache* create_cache();
 
-/** @brief Store a deep copy of matrix `m` under key `key`. */
-void cache_put(Cache* cache, const char* key, const Matrix* m);
+/** @brief Store matrix `m` under key `key`. The cache takes ownership of `m`,
+ *          and `m` will be freed when the cache entry is overwritten or the
+ * cache is cleared. The caller must ensure `m` is a deep copy if a distinct
+ * version is needed. */
+void cache_put(Cache* cache, const char* key, Matrix* m);
 
 /** @brief Retrieve a deep copy of a matrix by key, or NULL if not found. */
 Matrix* cache_get(Cache* cache, const char* key);
diff --git a/nn/include/feedforward.h b/nn/include/feedforward.h
index 37c91fd..6e62f65 100644
--- a/nn/include/feedforward.h
+++ b/nn/include/feedforward.h
@@ -26,10 +26,8 @@ void free_network(NeuralNetwork* nn);
 
 /**
  * @brief Run the forward pass and cache intermediates for backprop.
- *
  * @param nn Network pointer (non-NULL).
  * @param input Input matrix (batch_size x input_features).
- *              Shape must match the first layer's expected input shape.
  * @return Output activation of the last layer (batch_size x output_features).
  *         Caller owns and must free.
  */
diff --git a/nn/include/linalg.h b/nn/include/linalg.h
index d760990..b67e52b 100644
--- a/nn/include/linalg.h
+++ b/nn/include/linalg.h
@@ -40,20 +40,16 @@ Matrix* read_matrix(const char* filename);
 Matrix* create_matrix(size_t rows, size_t cols);
 /** @brief Deep copy a matrix. */
 Matrix* copy_matrix(const Matrix* m);
-/** @brief Flatten a matrix along an axis (implementation-specific). */
-Matrix* flatten_matrix(Matrix* m, int axis);
 /** @brief Fill all elements with a constant value. */
 void fill_matrix(Matrix* m, double n);
 /** @brief Fill with random values in an implementation-defined range. */
 void randomize_matrix(Matrix* m, double n);
 /** @brief Free a matrix and its data buffer. */
 void free_matrix(Matrix* m);
-/** @brief Write a matrix to a text file. */
-void write_matrix(Matrix* m, const char* filename);
 /** @brief Print a matrix to stdout (for debugging). */
 void print_matrix(Matrix* m);
 /** @brief Return the index of the maximum element (flattened argmax). */
-int matrix_argmax(Matrix* m);
+size_t matrix_argmax(Matrix* m);
 
 //============================
 // Matrix Operations
@@ -76,3 +72,9 @@ Matrix* dot_matrix(Matrix* a, Matrix* b);
 Matrix* transpose_matrix(Matrix* m);
 /** @brief Scale all elements by scalar n. */
 Matrix* scale_matrix(double n, Matrix* m);
+
+/** @brief Add a bias vector to each row of a matrix. */
+Matrix* add_bias_to_matrix(Matrix* m, Matrix* bias);
+
+/** @brief Sum the columns of a matrix, returning a row vector. */
+Matrix* sum_matrix_columns(Matrix* m);
diff --git a/nn/include/neural_network.h b/nn/include/neural_network.h
index 2e75f0f..4092442 100644
--- a/nn/include/neural_network.h
+++ b/nn/include/neural_network.h
@@ -19,11 +19,12 @@ typedef Matrix* (*ActivationFunc)(Matrix*);
 /**
  * @brief Fully connected layer parameters and activation.
  */
-typedef struct {
+typedef struct _Layer {
   Matrix* weights; /**< Weight matrix (D_in×D_out). */
   Matrix* bias;    /**< Bias vector as (1×D_out). */
 
-  ActivationFunc activation; /**< Activation function for this layer. */
+  activation_function
+      activation_type; /**< Type of activation function for this layer. */
 
   /** The leak parameter for Leaky ReLU activation. */
   double leak_parameter;
@@ -36,10 +37,26 @@ typedef struct {
 /**
  * @brief Neural network composed of sequential fully connected layers.
  */
-typedef struct {
+typedef struct _NeuralNetwork {
   Layer** layers;    /**< Array of layer pointers (length = num_layers). */
   size_t num_layers; /**< Number of layers. */
 
   /** Caches intermediate forward/backward values. */
   Cache* cache;
 } NeuralNetwork;
+
+/**
+ * @brief Prints a summary of the neural network's architecture.
+ * @param stream The output stream.
+ * @param nn The neural network to summarize.
+ */
+void fprint_network_summary(FILE* stream, const NeuralNetwork* nn);
+
+/**
+ * @brief Logs the training progress to a file.
+ * @param stream The output stream.
+ * @param epoch The current epoch.
+ * @param epochs The total number of epochs.
+ * @param loss The training loss.
+ */
+void flog_training_progress(FILE* stream, int epoch, int epochs, double loss);
diff --git a/nn/include/utils.h b/nn/include/utils.h
index 4ab29a0..ca06543 100644
--- a/nn/include/utils.h
+++ b/nn/include/utils.h
@@ -49,4 +49,4 @@ void log_message(LogLevel level, const char* format, ...);
 #define LOG_INFO(format, ...) log_message(LOG_LEVEL_INFO, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_message(LOG_LEVEL_WARN, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) \
-  log_message(LOG_LEVEL_ERROR, format, ##__VA_ARGS__)
\ No newline at end of file
+  log_message(LOG_LEVEL_ERROR, format, ##__VA_ARGS__)
diff --git a/nn/src/activation/activation.c b/nn/src/activation/activation.c
index 99ac2b4..a67acb2 100644
--- a/nn/src/activation/activation.c
+++ b/nn/src/activation/activation.c
@@ -20,6 +20,11 @@
 // Sigmoid Activation
 //============================
 
+/**
+ * @brief Applies the sigmoid activation function element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the sigmoid function applied to each element.
+ */
 Matrix* sigmoid(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying sigmoid activation to a %zux%zu matrix.", m->rows,
@@ -34,6 +39,12 @@ Matrix* sigmoid(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the sigmoid activation function
+ * element-wise to a matrix.
+ * @param m The input matrix (output of the sigmoid function).
+ * @return A new matrix with the sigmoid derivative applied to each element.
+ */
 Matrix* sigmoid_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying sigmoid_prime activation to a %zux%zu matrix.", m->rows,
@@ -53,6 +64,12 @@ Matrix* sigmoid_prime(Matrix* m) {
 // ReLU Activation
 //============================
 
+/**
+ * @brief Applies the ReLU (Rectified Linear Unit) activation function
+ * element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the ReLU function applied to each element.
+ */
 Matrix* relu(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying ReLU activation to a %zux%zu matrix.", m->rows, m->cols);
@@ -70,6 +87,12 @@ Matrix* relu(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the ReLU activation function element-wise
+ * to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the ReLU derivative applied to each element.
+ */
 Matrix* relu_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying ReLU_prime activation to a %zux%zu matrix.", m->rows,
@@ -92,6 +115,12 @@ Matrix* relu_prime(Matrix* m) {
 // Tanh Activation
 //============================
 
+/**
+ * @brief Applies the Hyperbolic Tangent (tanh) activation function element-wise
+ * to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the tanh function applied to each element.
+ */
 Matrix* tanh_activation(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Tanh activation to a %zux%zu matrix.", m->rows, m->cols);
@@ -105,6 +134,12 @@ Matrix* tanh_activation(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the Hyperbolic Tangent (tanh) activation
+ * function element-wise to a matrix.
+ * @param m The input matrix (output of the tanh function).
+ * @return A new matrix with the tanh derivative applied to each element.
+ */
 Matrix* tanh_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Tanh_prime activation to a %zux%zu matrix.", m->rows,
@@ -124,49 +159,17 @@ Matrix* tanh_prime(Matrix* m) {
 // Leaky ReLU Activation
 //============================
 
-// Without explicit definition, implemented in the functions below this
-// This will assume the alpha
-Matrix* leaky_relu(Matrix* m) {
-  ASSERT(m != NULL, "Input matrix is NULL.");
-  LOG_INFO("Applying Leaky ReLU activation to a %zux%zu matrix.", m->rows,
-           m->cols);
-
-  Matrix* result = create_matrix(m->rows, m->cols);
-  ASSERT(result != NULL, "Failed to create matrix.");
-  size_t total_elements = m->rows * m->cols;
-  for (size_t i = 0; i < total_elements; i++) {
-    if (m->matrix_data[i] > 0) {
-      result->matrix_data[i] = m->matrix_data[i];
-    } else {
-      result->matrix_data[i] = 0.01 * m->matrix_data[i];
-    }
-  }
-  return result;
-}
-
-Matrix* leaky_relu_prime(Matrix* m) {
-  ASSERT(m != NULL, "Input matrix is NULL.");
-  LOG_INFO("Applying Leaky ReLU_prime activation to a %zux%zu matrix.", m->rows,
-           m->cols);
-
-  Matrix* result = create_matrix(m->rows, m->cols);
-  ASSERT(result != NULL, "Failed to create matrix.");
-  size_t total_elements = m->rows * m->cols;
-  for (size_t i = 0; i < total_elements; i++) {
-    if (m->matrix_data[i] > 0) {
-      result->matrix_data[i] = 1;
-    } else {
-      result->matrix_data[i] = 0.01;
-    }
-  }
-  return result;
-}
-
-// For when users may require more explicit defintions of alpha
-Matrix* leaky_relu_with_alpha(Matrix* m, double leak_parameter) {
+/**
+ * @brief Applies the Leaky ReLU activation function element-wise to a matrix.
+ * @param m The input matrix.
+ * @param leak_parameter The leak parameter (alpha) for the Leaky ReLU. Must be
+ * non-negative.
+ * @return A new matrix with the Leaky ReLU function applied to each element.
+ */
+Matrix* leaky_relu(Matrix* m, double leak_parameter) {
   ASSERT(m != NULL, "Input matrix is NULL.");
-  // If I converted a non acceptable value of alpha into 0.01, it would bring in
-  // debug troubles.
+  // If I converted a non acceptable value of alpha into 0.01, it would bring
+  // in debug troubles.
   ASSERT(leak_parameter >= 0.0, "Alpha value must be non-negative.");
 
   LOG_INFO(
@@ -189,7 +192,15 @@ Matrix* leaky_relu_with_alpha(Matrix* m, double leak_parameter) {
   return result;
 }
 
-Matrix* leaky_relu_prime_with_alpha(Matrix* m, double leak_parameter) {
+/**
+ * @brief Computes the derivative of the Leaky ReLU activation function
+ * element-wise to a matrix.
+ * @param m The input matrix.
+ * @param leak_parameter The leak parameter (alpha) for the Leaky ReLU. Must be
+ * non-negative.
+ * @return A new matrix with the Leaky ReLU derivative applied to each element.
+ */
+Matrix* leaky_relu_prime(Matrix* m, double leak_parameter) {
   ASSERT(m != NULL, "Input matrix for leaky_relu_prime is NULL.");
   ASSERT(leak_parameter >= 0.0, "Alpha value must be non-negative.");
   LOG_INFO(
@@ -214,6 +225,11 @@ Matrix* leaky_relu_prime_with_alpha(Matrix* m, double leak_parameter) {
 // Sign Activation
 //============================
 
+/**
+ * @brief Applies the Sign activation function element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the Sign function applied to each element.
+ */
 Matrix* sign_activation(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Sign activation to a %zux%zu matrix.", m->rows, m->cols);
@@ -233,13 +249,22 @@ Matrix* sign_activation(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the Sign activation function element-wise
+ * to a matrix. The derivative of the sign function is 0 everywhere except at 0,
+ * where it is undefined. For backpropagation, the derivative is commonly
+ * approximated as 0.
+ * @param m The input matrix.
+ * @return A new matrix with the Sign derivative applied to each element (all
+ * zeros).
+ */
 Matrix* sign_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Sign_prime activation to a %zux%zu matrix.", m->rows,
            m->cols);
   // The derivative of the sign function is 0 everywhere except at 0, where it
-  // is undefined. For backpropagation, the derivative is commonly approximated
-  // as 0.
+  // is undefined. For backpropagation, the derivative is commonly
+  // approximated as 0.
   Matrix* result = create_matrix(m->rows, m->cols);
   ASSERT(result != NULL, "Failed to create matrix.");
   size_t total_elements = m->rows * m->cols;
@@ -253,6 +278,11 @@ Matrix* sign_prime(Matrix* m) {
 // Identity Activation
 //============================
 
+/**
+ * @brief Applies the Identity activation function element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix that is a copy of the input matrix.
+ */
 Matrix* identity_activation(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Identity activation to a %zux%zu matrix.", m->rows,
@@ -262,6 +292,13 @@ Matrix* identity_activation(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the Identity activation function
+ * element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix of the same size as the input, with all elements set
+ * to 1.0.
+ */
 Matrix* identity_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Identity_prime activation to a %zux%zu matrix.", m->rows,
@@ -280,6 +317,11 @@ Matrix* identity_prime(Matrix* m) {
 // Hard Tanh Activation
 //============================
 
+/**
+ * @brief Applies the Hard Tanh activation function element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the Hard Tanh function applied to each element.
+ */
 Matrix* hard_tanh(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Hard Tanh activation to a %zux%zu matrix.", m->rows,
@@ -300,6 +342,12 @@ Matrix* hard_tanh(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Computes the derivative of the Hard Tanh activation function
+ * element-wise to a matrix.
+ * @param m The input matrix.
+ * @return A new matrix with the Hard Tanh derivative applied to each element.
+ */
 Matrix* hard_tanh_prime(Matrix* m) {
   ASSERT(m != NULL, "Input matrix is NULL.");
   LOG_INFO("Applying Hard Tanh_prime activation to a %zux%zu matrix.", m->rows,
@@ -317,3 +365,81 @@ Matrix* hard_tanh_prime(Matrix* m) {
   }
   return result;
 }
+
+//============================
+// Softmax Activation
+//============================
+
+/**
+ * @brief Applies the Softmax activation function to a matrix.
+ * This function is typically used in the output layer of a neural network for
+ * multi-class classification. It normalizes the input values into a probability
+ * distribution.
+ * @param m The input matrix.
+ * @return A new matrix with the Softmax function applied.
+ */
+Matrix* softmax(Matrix* m) {
+  ASSERT(m != NULL, "Input matrix is NULL.");
+  LOG_INFO("Applying softmax activation to a %zux%zu matrix.", m->rows,
+           m->cols);
+
+  Matrix* result = create_matrix(m->rows, m->cols);
+  ASSERT(result != NULL, "Failed to create matrix.");
+
+  for (size_t i = 0; i < m->rows; i++) {
+    double max_val = m->matrix_data[i * m->cols];
+    for (size_t j = 1; j < m->cols; j++) {
+      if (m->matrix_data[i * m->cols + j] > max_val) {
+        max_val = m->matrix_data[i * m->cols + j];
+      }
+    }
+
+    double sum = 0.0;
+    for (size_t j = 0; j < m->cols; j++) {
+      sum += exp(m->matrix_data[i * m->cols + j] - max_val);
+    }
+
+    for (size_t j = 0; j < m->cols; j++) {
+      result->matrix_data[i * m->cols + j] =
+          exp(m->matrix_data[i * m->cols + j] - max_val) / sum;
+    }
+  }
+  return result;
+}
+
+/**
+ * @brief Computes the derivative of the Softmax activation function
+ * element-wise to a matrix. This is typically used in conjunction with a loss
+ * function like cross-entropy where the derivative simplifies to output * (1 -
+ * output).
+ * @param m The input matrix (output of the softmax function).
+ * @return A new matrix with the Softmax derivative applied to each element.
+ */
+Matrix* softmax_prime(Matrix* m) {
+  ASSERT(m != NULL, "Input matrix is NULL.");
+  LOG_INFO("Applying softmax_prime activation to a %zux%zu matrix.", m->rows,
+           m->cols);
+
+  Matrix* result = create_matrix(m->rows, m->cols);
+  ASSERT(result != NULL, "Failed to create matrix.");
+  size_t total_elements = m->rows * m->cols;
+  for (size_t i = 0; i < total_elements; i++) {
+    result->matrix_data[i] = m->matrix_data[i] * (1.0 - m->matrix_data[i]);
+  }
+  return result;
+}
+
+const char* activation_to_string(activation_function func) {
+  switch (func) {
+    case SIGMOID:
+      return "SIGMOID";
+    case RELU:
+      return "RELU";
+    case LEAKY_RELU:
+      return "LEAKY_RELU";
+    case SOFTMAX:
+      return "SOFTMAX";
+    default:
+      return "UNKNOWN";
+  }
+}
diff --git a/nn/src/cache/cache.c b/nn/src/cache/cache.c
index ebdffb1..6632e89 100644
--- a/nn/src/cache/cache.c
+++ b/nn/src/cache/cache.c
@@ -11,18 +11,29 @@
 #include "linalg.h"
 
 // A linked list is used to handle collisions at each bucket.
+/**
+ * @brief Represents an entry in the cache, storing a key-matrix pair.
+ * Uses a linked list for collision resolution in the hash map.
+ */
 typedef struct CacheEntry {
-  char* key;
-  Matrix* m;
-  struct CacheEntry* next;
+  char* key; /**< The key associated with the matrix. */
+  Matrix* m; /**< The matrix stored in this entry. */
+  struct CacheEntry*
+      next; /**< Pointer to the next entry in case of collision. */
 } CacheEntry;
 
 #define HASH_MAP_SIZE 1024
 
 struct Cache {
-  CacheEntry* entries[HASH_MAP_SIZE];
+  CacheEntry* entries[HASH_MAP_SIZE]; /**< Array of pointers to CacheEntry,
+                                         forming the hash table buckets. */
 };
 
+/**
+ * @brief Computes a hash value for a given string key.
+ * @param key The string key to hash.
+ * @return An unsigned integer hash value, modulo HASH_MAP_SIZE.
+ */
 static unsigned int hash(const char* key) {
   // This hash function can be manipulated. You could input a key
   // that could overflow a 32 bit uint, so I've made it a 64 bit
@@ -42,6 +53,11 @@ static unsigned int hash(const char* key) {
 // Cache Functions
 //------------------------------
 
+/**
+ * @brief Creates and initializes a new, empty cache.
+ * @return A pointer to the newly created Cache, or NULL if memory allocation
+ * fails.
+ */
 Cache* create_cache() {
   Cache* cache = (Cache*)malloc(sizeof(Cache));
   if (cache == NULL) {
@@ -54,7 +70,15 @@ Cache* create_cache() {
   return cache;
 }
 
-void cache_put(Cache* cache, const char* key, const Matrix* m) {
+/**
+ * @brief Stores a matrix in the cache under a specified key.
+ * If the key already exists, the old matrix is freed and replaced with the new
+ * one. The cache takes ownership of the matrix `m`.
+ * @param cache A pointer to the Cache structure.
+ * @param key The string key for the matrix.
+ * @param m A pointer to the Matrix to be stored.
+ */
+void cache_put(Cache* cache, const char* key, Matrix* m) {
   if (cache == NULL || key == NULL || m == NULL) {
     return;
   }
@@ -66,7 +90,7 @@ void cache_put(Cache* cache, const char* key, const Matrix* m) {
     if (strcmp(current->key, key) == 0) {
       // Key found, free existing matrix and update with new one.
       free_matrix(current->m);
-      current->m = copy_matrix(m);  // Deep copy
+      current->m = m;
       return;
     }
     current = current->next;
@@ -77,11 +101,18 @@ void cache_put(Cache* cache, const char* key, const Matrix* m) {
     return;
   }
   new_entry->key = strdup(key);
-  new_entry->m = copy_matrix(m);  // Deep copy to prevent side effects
+  new_entry->m = m;
   new_entry->next = cache->entries[index];
   cache->entries[index] = new_entry;
 }
 
+/**
+ * @brief Retrieves a deep copy of a matrix from the cache using its key.
+ * @param cache A pointer to the Cache structure.
+ * @param key The string key of the matrix to retrieve.
+ * @return A deep copy of the stored Matrix, or NULL if the key is not found or
+ * allocation fails. The caller is responsible for freeing the returned matrix.
+ */
 Matrix* cache_get(Cache* cache, const char* key) {
   if (cache == NULL || key == NULL) {
     return NULL;
@@ -97,6 +128,11 @@ Matrix* cache_get(Cache* cache, const char* key) {
   return NULL;
 }
 
+/**
+ * @brief Clears all entries from the cache, freeing associated memory for keys
+ * and matrices. The cache structure itself is not freed.
+ * @param cache A pointer to the Cache structure to clear.
+ */
 void clear_cache(Cache* cache) {
   if (cache == NULL) {
     return;
@@ -114,6 +150,10 @@ void clear_cache(Cache* cache) {
   }
 }
 
+/**
+ * @brief Frees all entries in the cache and the cache structure itself.
+ * @param cache A pointer to the Cache structure to free.
+ */
 void free_cache(Cache* cache) {
   if (cache == NULL) {
     return;
diff --git a/nn/src/linalg/io.c b/nn/src/linalg/io.c
index 77d4142..36194b2 100644
--- a/nn/src/linalg/io.c
+++ b/nn/src/linalg/io.c
@@ -15,6 +15,13 @@
 // Functions for Matrix IO
 //============================
 
+/**
+ * @brief Reads a matrix from a text file.
+ * The file format is expected to be: first line for rows, second for columns,
+ * then matrix data.
+ * @param filename The path to the file to read.
+ * @return A pointer to the newly created Matrix, or NULL if an error occurs.
+ */
 Matrix* read_matrix(const char* filename) {
   LOG_INFO("Attempting to load matrix from file: %s", filename);
 
@@ -24,23 +31,39 @@ Matrix* read_matrix(const char* filename) {
   char entry[1024];
   size_t rows = 0, cols = 0;
 
+  char* endptr;
+  Matrix* m = NULL;  // Initialize m to NULL
+
   if (fgets(entry, sizeof(entry), file) == NULL) {
     LOG_ERROR("Could not read rows from file: %s", filename);
     fclose(file);
     return NULL;
   }
-  rows = atoi(entry);
+  rows = strtol(entry, &endptr, 10);
+  if (endptr == entry ||
+      (*endptr != '\n' && *endptr != '\0')) {  // Added parentheses
+    LOG_ERROR("Invalid row format in file: %s", filename);
+    fclose(file);
+    return NULL;
+  }
 
   if (fgets(entry, sizeof(entry), file) == NULL) {
     LOG_ERROR("Could not read columns from file: %s", filename);
     fclose(file);
     return NULL;
   }
-  cols = atoi(entry);
+  cols = strtol(entry, &endptr, 10);
+  if (endptr == entry ||
+      (*endptr != '\n' && *endptr != '\0')) {  // Added parentheses
+    LOG_ERROR("Invalid column format in file: %s", filename);
+    // No need to free m here, as it's still NULL
+    fclose(file);
+    return NULL;
+  }
 
   ASSERT(rows > 0 && cols > 0, "Invalid matrix dimensions read from file.");
 
-  Matrix* m = create_matrix(rows, cols);
+  m = create_matrix(rows, cols);  // Assign to m after successful creation
 
   for (size_t i = 0; i < rows; i++) {
     if (fgets(entry, sizeof(entry), file) == NULL) {
@@ -51,8 +74,17 @@ Matrix* read_matrix(const char* filename) {
     }
 
     char* line_ptr = entry;
+    char* prev_line_ptr = entry;
     for (size_t j = 0; j < cols; j++) {
       m->matrix_data[i * cols + j] = strtod(line_ptr, &line_ptr);
+      if (line_ptr == prev_line_ptr) {
+        LOG_ERROR("Invalid number format in matrix data at row %zu, col %zu.",
+                  i, j);
+        free_matrix(m);
+        fclose(file);
+        return NULL;
+      }
+      prev_line_ptr = line_ptr;
     }
   }
 
@@ -80,6 +112,11 @@ Matrix* create_matrix(size_t rows, size_t cols) {
   return matrix;
 }
 
+/**
+ * @brief Creates a deep copy of an existing matrix.
+ * @param m A pointer to the source Matrix to be copied.
+ * @return A pointer to the newly created deep copy of the matrix.
+ */
 Matrix* copy_matrix(const Matrix* m) {
   ASSERT(m != NULL, "Input matrix for copy is NULL.");
 
@@ -94,39 +131,11 @@ Matrix* copy_matrix(const Matrix* m) {
   return new_matrix;
 }
 
-// Helper function for flattening
-Matrix* flatten_column_wise(const Matrix* m) {
-  Matrix* new_matrix = create_matrix(m->rows * m->cols, 1);
-
-  size_t k = 0;
-  for (size_t j = 0; j < m->cols; j++) {
-    for (size_t i = 0; i < m->rows; i++) {
-      new_matrix->matrix_data[k] = m->matrix_data[i * m->cols + j];
-      k++;
-    }
-  }
-  return new_matrix;
-}
-
-Matrix* flatten_matrix(Matrix* m, int axis) {
-  ASSERT(m != NULL, "Input matrix to flatten is NULL.");
-  ASSERT(axis == 0 || axis == 1,
-         "Axis must be 0 (row-wise) or 1 (column-wise).");
-
-  if (axis == 0) {
-    LOG_INFO(
-        "Flattening matrix row-wise. No operation needed as data is "
-        "already "
-        "contiguous.");
-    m->cols = m->rows * m->cols;
-    m->rows = 1;
-    return m;
-  } else {
-    LOG_INFO("Flattening matrix column-wise. A new matrix will be created.");
-    return flatten_column_wise(m);
-  }
-}
-
+/**
+ * @brief Fills all elements of a matrix with a specified scalar value.
+ * @param m A pointer to the Matrix to be filled.
+ * @param n The double value to fill the matrix with.
+ */
 void fill_matrix(Matrix* m, double n) {
   ASSERT(m != NULL, "Input matrix for fill_matrix is NULL.");
   LOG_INFO("Filling a %zux%zu matrix with the value %.2f.", m->rows, m->cols,
@@ -139,6 +148,13 @@ void fill_matrix(Matrix* m, double n) {
   }
 }
 
+/**
+ * @brief Randomizes the elements of a matrix within a specific range.
+ * The range is determined by `n` (typically the number of input features) to
+ * help prevent vanishing/exploding gradients.
+ * @param m A pointer to the Matrix to be randomized.
+ * @param n A scaling factor used to determine the range of random values.
+ */
 void randomize_matrix(Matrix* m, double n) {
   LOG_INFO("Randomizing a %zux%zu matrix.", m->rows, m->cols);
   // Apparently a 1/n or 1/n^2 scaling leads to a vanishing gradient problem
@@ -155,6 +171,10 @@ void randomize_matrix(Matrix* m, double n) {
   LOG_INFO("Matrix randomized successfully.");
 }
 
+/**
+ * @brief Frees the memory allocated for a matrix.
+ * @param m A pointer to the Matrix to be freed.
+ */
 void free_matrix(Matrix* m) {
   LOG_INFO("Freeing matrix at address %p.", m);
   if (m == NULL) {
@@ -168,6 +188,11 @@ void free_matrix(Matrix* m) {
   LOG_INFO("Matrix freed successfully.");
 }
 
+/**
+ * @brief Prints the elements of a matrix to standard output for debugging
+ * purposes.
+ * @param m A pointer to the Matrix to be printed.
+ */
 void print_matrix(Matrix* m) {
   ASSERT(m != NULL, "Input matrix for print is NULL.");
   LOG_INFO("Printing matrix of size %zux%zu.", m->rows, m->cols);
@@ -179,6 +204,13 @@ void print_matrix(Matrix* m) {
   }
 }
 
+/**
+ * @brief Writes a matrix to a text file.
+ * The format written is: rows\n, cols\n, then matrix data with space-separated
+ * values.
+ * @param m A pointer to the Matrix to be written.
+ * @param filename The path to the file where the matrix will be saved.
+ */
 void write_matrix(Matrix* m, const char* filename) {
   ASSERT(m != NULL, "Input matrix for save is NULL.");
   LOG_INFO("Saving a %zux%zu matrix to file: %s", m->rows, m->cols, filename);
@@ -200,19 +232,24 @@ void write_matrix(Matrix* m, const char* filename) {
   LOG_INFO("Matrix saved successfully.");
 }
 
-int matrix_argmax(Matrix* m) {
+/**
+ * @brief Finds the index of the maximum element in a flattened matrix.
+ * @param m A pointer to the Matrix.
+ * @return The 0-based index of the maximum element.
+ */
+size_t matrix_argmax(Matrix* m) {
   ASSERT(m != NULL, "Input matrix for argmax is NULL.");
-  ASSERT(m->cols == 1, "Input must be a column vector (Mx1).");
 
-  double maxValue = INT_MIN;
-  int maxIndex = 0;
+  double maxValue = m->matrix_data[0];
+  size_t maxIndex = 0;
 
-  for (size_t i = 0; i < m->rows; i++) {
+  size_t total_elements = m->rows * m->cols;
+  for (size_t i = 1; i < total_elements; i++) {
     if (m->matrix_data[i] > maxValue) {
-      maxIndex = (int)i;
+      maxIndex = i;
       maxValue = m->matrix_data[i];
     }
   }
-  LOG_INFO("Max value found at index %d.", maxIndex);
+  LOG_INFO("Max value found at index %zu.", maxIndex);
   return maxIndex;
-}
+}
\ No newline at end of file
diff --git a/nn/src/linalg/operations.c b/nn/src/linalg/operations.c
index 99a0590..8c66b71 100644
--- a/nn/src/linalg/operations.c
+++ b/nn/src/linalg/operations.c
@@ -19,6 +19,11 @@
 // Functions for Matrix Operations
 //============================
 
+/**
+ * @brief Creates an identity matrix of size n x n.
+ * @param n The dimension of the square identity matrix.
+ * @return A pointer to the newly created identity Matrix.
+ */
 Matrix* identity_matrix(size_t n) {
   LOG_INFO("Creating a %zux%zu identity matrix.", n, n);
   ASSERT(n > 0, "Matrix size must be greater than 0.");
@@ -49,6 +54,12 @@ Matrix* add_matrix(Matrix* a, Matrix* b) {
   return result;
 }
 
+/**
+ * @brief Performs element-wise subtraction of two matrices.
+ * @param a The first matrix (minuend).
+ * @param b The second matrix (subtrahend).
+ * @return A new matrix containing the result of a - b.
+ */
 Matrix* subtract_matrix(Matrix* a, Matrix* b) {
   ASSERT(a != NULL && b != NULL, "Input matrices cannot be NULL.");
   ASSERT(a->rows == b->rows && a->cols == b->cols,
@@ -66,6 +77,13 @@ Matrix* subtract_matrix(Matrix* a, Matrix* b) {
   return result;
 }
 
+/**
+ * @brief Performs element-wise multiplication of two matrices (Hadamard
+ * product).
+ * @param a The first matrix.
+ * @param b The second matrix.
+ * @return A new matrix containing the element-wise product of a and b.
+ */
 Matrix* multiply_matrix(Matrix* a, Matrix* b) {
   ASSERT(a != NULL && b != NULL, "Input matrices cannot be NULL.");
   ASSERT(a->rows == b->rows && a->cols == b->cols,
@@ -84,6 +102,12 @@ Matrix* multiply_matrix(Matrix* a, Matrix* b) {
   return result;
 }
 
+/**
+ * @brief Applies a given function to each element of a matrix.
+ * @param func A function pointer that takes a double and returns a double.
+ * @param m The input matrix.
+ * @return A new matrix with the function applied to each element.
+ */
 Matrix* apply_onto_matrix(double (*func)(double), Matrix* m) {
   ASSERT(m != NULL, "Input matrix cannot be NULL.");
   LOG_INFO("Applying a function to each element of a %zux%zu matrix.", m->rows,
@@ -99,6 +123,12 @@ Matrix* apply_onto_matrix(double (*func)(double), Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Adds a scalar value to each element of a matrix.
+ * @param m The input matrix.
+ * @param n The scalar value to add.
+ * @return A new matrix with the scalar added to each element.
+ */
 Matrix* add_scalar_to_matrix(Matrix* m, double n) {
   ASSERT(m != NULL, "Input matrix cannot be NULL.");
   LOG_INFO("Adding scalar %.2f to a %zux%zu matrix.", n, m->rows, m->cols);
@@ -113,6 +143,12 @@ Matrix* add_scalar_to_matrix(Matrix* m, double n) {
   return result;
 }
 
+/**
+ * @brief Performs the dot product (matrix multiplication) of two matrices.
+ * @param a The first matrix.
+ * @param b The second matrix.
+ * @return A new matrix containing the result of the dot product a * b.
+ */
 Matrix* dot_matrix(Matrix* a, Matrix* b) {
   ASSERT(a != NULL && b != NULL, "Input matrices cannot be NULL.");
   ASSERT(a->cols == b->rows,
@@ -139,6 +175,11 @@ Matrix* dot_matrix(Matrix* a, Matrix* b) {
   return result;
 }
 
+/**
+ * @brief Transposes a matrix.
+ * @param m The input matrix.
+ * @return A new matrix that is the transpose of the input matrix.
+ */
 Matrix* transpose_matrix(Matrix* m) {
   ASSERT(m != NULL, "Input matrix cannot be NULL.");
   LOG_INFO("Transposing a %zux%zu matrix.", m->rows, m->cols);
@@ -156,6 +197,12 @@ Matrix* transpose_matrix(Matrix* m) {
   return result;
 }
 
+/**
+ * @brief Scales all elements of a matrix by a scalar value.
+ * @param n The scalar value to multiply by.
+ * @param m The input matrix.
+ * @return A new matrix with all elements scaled by n.
+ */
 Matrix* scale_matrix(double n, Matrix* m) {
   ASSERT(m != NULL, "Input matrix cannot be NULL.");
   LOG_INFO("Scaling a %zux%zu matrix by %.2f.", m->rows, m->cols, n);
@@ -169,3 +216,52 @@ Matrix* scale_matrix(double n, Matrix* m) {
   LOG_INFO("Matrix scaling complete.");
   return result;
 }
+
+/**
+ * @brief Adds a bias vector to each row of a matrix.
+ * @param m The input matrix.
+ * @param bias The bias vector (must be a 1xN row vector where N is m->cols).
+ * @return A new matrix with the bias added to each row.
+ */
+Matrix* add_bias_to_matrix(Matrix* m, Matrix* bias) {
+  ASSERT(m != NULL, "Input matrix is NULL.");
+  ASSERT(bias != NULL, "Bias matrix is NULL.");
+  ASSERT(bias->rows == 1, "Bias must be a row vector.");
+  ASSERT(m->cols == bias->cols,
+         "Matrix and bias dimensions are incompatible for addition.");
+
+  Matrix* result = create_matrix(m->rows, m->cols);
+  ASSERT(result != NULL, "Failed to create matrix for bias addition.");
+
+  for (size_t i = 0; i < m->rows; i++) {
+    for (size_t j = 0; j < m->cols; j++) {
+      result->matrix_data[i * m->cols + j] =
+          m->matrix_data[i * m->cols + j] + bias->matrix_data[j];
+    }
+  }
+
+  return result;
+}
+
+/**
+ * @brief Sums the columns of a matrix, returning a row vector.
+ * @param m The input matrix.
+ * @return A new 1xN matrix (row vector) where each element is the sum of the
+ * corresponding column in m.
+ */
+Matrix* sum_matrix_columns(Matrix* m) {
+  ASSERT(m != NULL, "Input matrix is NULL.");
+
+  Matrix* result = create_matrix(1, m->cols);
+  ASSERT(result != NULL, "Failed to create matrix for column summation.");
+
+  for (size_t j = 0; j < m->cols; j++) {
+    double sum = 0;
+    for (size_t i = 0; i < m->rows; i++) {
+      sum += m->matrix_data[i * m->cols + j];
+    }
+    result->matrix_data[j] = sum;
+  }
+
+  return result;
+}
diff --git a/nn/src/neural_network/backprop.c b/nn/src/neural_network/backprop.c
index 5da039c..48f7f55 100644
--- a/nn/src/neural_network/backprop.c
+++ b/nn/src/neural_network/backprop.c
@@ -5,6 +5,7 @@
 #include "backprop.h"
 
 #include <stddef.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 
@@ -16,48 +17,52 @@
 
 // Select and compute activation derivative for a layer given its pre-activation
 // input z. Handles common activations and leaky ReLU with optional alpha.
+/**
+ * @brief Selects and computes the derivative of the activation function for a
+ * given layer.
+ * @param layer A pointer to the Layer structure containing the activation type
+ * and parameters.
+ * @param z A pointer to the pre-activation matrix (input to the activation
+ * function).
+ * @return A new matrix containing the element-wise derivative of the activation
+ * function applied to z.
+ */
 static Matrix* activation_derivative_for_layer(const Layer* layer, Matrix* z) {
   ASSERT(layer != NULL, "Layer cannot be NULL.");
   ASSERT(z != NULL, "Pre-activation matrix z cannot be NULL.");
 
-  if (layer->activation == sigmoid) {
-    return sigmoid_prime(z);
-  }
-  if (layer->activation == relu) {
-    return relu_prime(z);
-  }
-  if (layer->activation == tanh_activation) {
-    return tanh_prime(z);
-  }
-  if (layer->activation == leaky_relu) {
-    return leaky_relu_prime(z);
+  switch (layer->activation_type) {
+    case SIGMOID:
+      return sigmoid_prime(z);
+    case RELU:
+      return relu_prime(z);
+    case TANH:
+      return tanh_prime(z);
+    case LEAKY_RELU:
+      return leaky_relu_prime(z, layer->leak_parameter);
+    case SIGN:
+      return sign_prime(z);
+    case IDENTITY:
+      return identity_prime(z);
+    case HARD_TANH:
+      return hard_tanh_prime(z);
+    default:
+      LOG_WARN(
+          "Unknown activation function, defaulting derivative to identity.");
+      return identity_prime(z);
   }
-  if (layer->activation == leaky_relu_with_alpha) {
-    return leaky_relu_prime_with_alpha(z, layer->leak_parameter);
-  }
-  if (layer->activation == sign_activation) {
-    return sign_prime(z);
-  }
-  if (layer->activation == identity_activation) {
-    return identity_prime(z);
-  }
-  if (layer->activation == hard_tanh) {
-    return hard_tanh_prime(z);
-  }
-
-  // Default: identity derivative
-  LOG_WARN("Unknown activation function, defaulting derivative to identity.");
-  return identity_prime(z);
 }
 
 void backpropagate(NeuralNetwork* nn, const Matrix* y_true,
-                   LossFunction loss_func, LossFunctionGrad loss_func_grad) {
+                   LossFunctionType loss_type,
+                   LossFunctionGrad loss_func_grad) {
   ASSERT(nn != NULL, "Neural Network pointer cannot be NULL.");
   ASSERT(nn->cache != NULL, "Cache cannot be NULL.");
   ASSERT(y_true != NULL, "Ground truth matrix cannot be NULL.");
   ASSERT(loss_func_grad != NULL, "Loss gradient function cannot be NULL.");
 
   size_t last_index = nn->num_layers - 1;
+  Layer* last_layer = nn->layers[last_index];
 
   // Get y_hat from cache (activation of last layer)
   char a_last_key[32];
@@ -65,20 +70,30 @@ void backpropagate(NeuralNetwork* nn, const Matrix* y_true,
   Matrix* y_hat = cache_get(nn->cache, a_last_key);
   ASSERT(y_hat != NULL, "Cached prediction (y_hat) not found.");
 
-  // dL/da for output layer
-  Matrix* dL_da = loss_func_grad(y_hat, y_true);
-  ASSERT(dL_da != NULL, "Loss gradient returned NULL.");
-
-  // delta for output layer: dL/dz = dL/da .* a'(z)
-  char z_last_key[32];
-  sprintf(z_last_key, "z_%zu", last_index);
-  Matrix* z_last = cache_get(nn->cache, z_last_key);
-  ASSERT(z_last != NULL, "Cached z for last layer not found.");
-
-  Matrix* act_prime_last =
-      activation_derivative_for_layer(nn->layers[last_index], z_last);
-  Matrix* delta_last = multiply_matrix(dL_da, act_prime_last);
-  ASSERT(delta_last != NULL, "Failed to compute delta for last layer.");
+  Matrix* delta_last;
+  // Special case for Softmax with CCE
+  if (last_layer->activation_type == SOFTMAX && loss_type == CCE) {
+    delta_last = subtract_matrix(y_hat, (Matrix*)y_true);
+  } else {
+    // dL/da for output layer
+    Matrix* dL_da = loss_func_grad(y_hat, y_true);
+    ASSERT(dL_da != NULL, "Loss gradient returned NULL.");
+
+    // delta for output layer: dL/dz = dL/da .* a'(z)
+    char z_last_key[32];
+    sprintf(z_last_key, "z_%zu", last_index);
+    Matrix* z_last = cache_get(nn->cache, z_last_key);
+    ASSERT(z_last != NULL, "Cached z for last layer not found.");
+
+    Matrix* act_prime_last =
+        activation_derivative_for_layer(last_layer, z_last);
+    delta_last = multiply_matrix(dL_da, act_prime_last);
+    ASSERT(delta_last != NULL, "Failed to compute delta for last layer.");
+
+    free_matrix(dL_da);
+    free_matrix(z_last);
+    free_matrix(act_prime_last);
+  }
 
   char delta_last_key[32];
   sprintf(delta_last_key, "delta_%zu", last_index);
@@ -86,10 +101,6 @@ void backpropagate(NeuralNetwork* nn, const Matrix* y_true,
 
   // Clean up temporaries for last layer
   free_matrix(y_hat);
-  free_matrix(dL_da);
-  free_matrix(z_last);
-  free_matrix(act_prime_last);
-  free_matrix(delta_last);
 
   // Backpropagate through hidden layers
   for (size_t i = last_index - 1; i != SIZE_MAX; i--) {
@@ -124,10 +135,18 @@ void backpropagate(NeuralNetwork* nn, const Matrix* y_true,
     free_matrix(propagated);
     free_matrix(z_i);
     free_matrix(act_prime_i);
-    free_matrix(delta_i);
   }
 }
 
+/**
+ * @brief Calculates the gradient of the weights for a specific layer during
+ * backpropagation.
+ * @param cache A pointer to the Cache containing intermediate values.
+ * @param layer_index The index of the current layer.
+ * @param total_layers The total number of layers in the neural network.
+ * @return A new matrix representing the gradient of the weights for the
+ * specified layer.
+ */
 Matrix* calculate_weight_gradient(const Cache* cache, size_t layer_index,
                                   size_t total_layers) {
   ASSERT(cache != NULL, "Cache cannot be NULL.");
@@ -161,6 +180,15 @@ Matrix* calculate_weight_gradient(const Cache* cache, size_t layer_index,
   return grad_W;
 }
 
+/**
+ * @brief Calculates the gradient of the biases for a specific layer during
+ * backpropagation.
+ * @param cache A pointer to the Cache containing intermediate values.
+ * @param layer_index The index of the current layer.
+ * @param total_layers The total number of layers in the neural network.
+ * @return A new matrix representing the gradient of the biases for the
+ * specified layer.
+ */
 Matrix* calculate_bias_gradient(const Cache* cache, size_t layer_index,
                                 size_t total_layers) {
   ASSERT(cache != NULL, "Cache cannot be NULL.");
@@ -171,6 +199,8 @@ Matrix* calculate_bias_gradient(const Cache* cache, size_t layer_index,
   Matrix* delta_i = cache_get((Cache*)cache, delta_key);
   ASSERT(delta_i != NULL, "Cached delta for layer not found.");
 
-  // For single-sample case, bias gradient equals delta
-  return delta_i;  // already a deep copy from cache
+  Matrix* db = sum_matrix_columns(delta_i);
+  free_matrix(delta_i);
+
+  return db;
 }
diff --git a/nn/src/neural_network/feedforward.c b/nn/src/neural_network/feedforward.c
index 3d36183..7cd105c 100644
--- a/nn/src/neural_network/feedforward.c
+++ b/nn/src/neural_network/feedforward.c
@@ -8,6 +8,7 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "activation.h"
 #include "linalg.h"
 #include "neural_network.h"
 #include "utils.h"
@@ -25,6 +26,10 @@ NeuralNetwork* create_network(size_t num_layers) {
     free(nn);
     return NULL;
   }
+  // Initialize layer pointers to NULL
+  for (size_t i = 0; i < num_layers; i++) {
+    nn->layers[i] = NULL;
+  }
 
   nn->num_layers = num_layers;
   nn->cache = create_cache();
@@ -37,6 +42,11 @@ NeuralNetwork* create_network(size_t num_layers) {
   return nn;
 }
 
+/**
+ * @brief Frees all memory associated with a neural network.
+ * This includes layers, weights, biases, and the cache.
+ * @param nn A pointer to the NeuralNetwork structure to be freed.
+ */
 void free_network(NeuralNetwork* nn) {
   if (nn == NULL) {
     return;
@@ -62,16 +72,25 @@ void free_network(NeuralNetwork* nn) {
   free(nn);
 }
 
+/**
+ * @brief Performs a forward pass through the neural network.
+ * Computes the output of the network for a given input and caches intermediate
+ * values.
+ * @param nn A pointer to the NeuralNetwork structure.
+ * @param input A pointer to the input Matrix.
+ * @return A new matrix containing the output of the last layer of the network.
+ * The caller is responsible for freeing this matrix.
+ */
 Matrix* feedforward(NeuralNetwork* nn, const Matrix* input) {
   ASSERT(nn != NULL, "Neural Network pointer cannot be NULL.");
   ASSERT(input != NULL, "Input matrix cannot be NULL.");
-  ASSERT(input->rows == nn->layers[0]->weights->cols,
+  ASSERT(input->cols == nn->layers[0]->weights->rows,
          "Input dimensions must match network dimensions.");
 
   Matrix* current_output = copy_matrix(input);
   ASSERT(current_output != NULL, "Failed to copy input matrix.");
 
-  cache_put(nn->cache, "input", current_output);
+  cache_put(nn->cache, "input", copy_matrix(current_output));
 
   for (size_t i = 0; i < nn->num_layers; i++) {
     Layer* current_layer = nn->layers[i];
@@ -84,7 +103,7 @@ Matrix* feedforward(NeuralNetwork* nn, const Matrix* input) {
                z_linear->cols == current_layer->weights->cols,
            "Unexpected shape from dot product.");
 
-    Matrix* z = add_matrix(z_linear, current_layer->bias);
+    Matrix* z = add_bias_to_matrix(z_linear, current_layer->bias);
     ASSERT(z != NULL, "Bias add failed.");
     ASSERT(z->rows == z_linear->rows && z->cols == z_linear->cols,
            "Unexpected shape from bias add.");
@@ -92,14 +111,38 @@ Matrix* feedforward(NeuralNetwork* nn, const Matrix* input) {
     // Cache the intermediate pre-activation value (z).
     char z_key[32];
     sprintf(z_key, "z_%zu", i);
-    cache_put(nn->cache, z_key, z);
+    cache_put(nn->cache, z_key, copy_matrix(z));
 
     Matrix* a = NULL;
-    // Handle special-case leaky ReLU with custom alpha.
-    if (current_layer->activation == leaky_relu_with_alpha) {
-      a = leaky_relu_with_alpha(z, current_layer->leak_parameter);
-    } else {
-      a = current_layer->activation(z);
+    switch (current_layer->activation_type) {
+      case SIGMOID:
+        a = sigmoid(z);
+        break;
+      case RELU:
+        a = relu(z);
+        break;
+      case TANH:
+        a = tanh_activation(z);
+        break;
+      case LEAKY_RELU:
+        a = leaky_relu(z, current_layer->leak_parameter);
+        break;
+      case SIGN:
+        a = sign_activation(z);
+        break;
+      case IDENTITY:
+        a = identity_activation(z);
+        break;
+      case HARD_TANH:
+        a = hard_tanh(z);
+        break;
+      case SOFTMAX:
+        a = softmax(z);
+        break;
+      default:
+        LOG_WARN("Unknown activation function, defaulting to identity.");
+        a = identity_activation(z);
+        break;
     }
     ASSERT(a != NULL, "Activation failed.");
     ASSERT(a->rows == z->rows && a->cols == z->cols,
@@ -107,7 +150,7 @@ Matrix* feedforward(NeuralNetwork* nn, const Matrix* input) {
 
     char a_key[32];
     sprintf(a_key, "a_%zu", i);
-    cache_put(nn->cache, a_key, a);
+    cache_put(nn->cache, a_key, copy_matrix(a));
 
     free_matrix(z_linear);
     free_matrix(z);
diff --git a/nn/src/utils/utils.c b/nn/src/utils/utils.c
new file mode 100644
index 0000000..465e7df
--- /dev/null
+++ b/nn/src/utils/utils.c
@@ -0,0 +1,59 @@
+/**
+ * @file utils.c
+ * @brief Implementation of logging and assertion utilities.
+ */
+
+#include "utils.h"
+
+#include <stdarg.h>
+#include <string.h>
+
+/**
+ * @brief Returns the string representation of a given LogLevel.
+ * @param level The LogLevel enum value.
+ * @return A string literal representing the log level.
+ */
+const char* get_log_level_string(LogLevel level) {
+  switch (level) {
+    case LOG_LEVEL_DEBUG:
+      return "DEBUG";
+    case LOG_LEVEL_INFO:
+      return "INFO";
+    case LOG_LEVEL_WARN:
+      return "WARN";
+    case LOG_LEVEL_ERROR:
+      return "ERROR";
+    default:
+      return "UNKNOWN";
+  }
+}
+
+/**
+ * @brief Logs a formatted message to stdout or stderr based on the log level.
+ * Messages with level WARN or ERROR are directed to stderr.
+ * @param level The LogLevel of the message.
+ * @param format The format string for the message.
+ * @param ... Variable arguments to be formatted according to 'format'.
+ */
+void log_message(LogLevel level, const char* format, ...) {
+  if (level < MIN_LOG_LEVEL) {
+    return;
+  }
+
+  // include time
+  time_t now = time(NULL);
+  const struct tm* t = localtime(&now);
+  char time_str[20];
+  strftime(time_str, sizeof(time_str), "%Y-%m-%d %H:%M:%S", t);
+
+  FILE* stream = (level >= LOG_LEVEL_WARN) ? stderr : stdout;
+
+  va_list args;
+  va_start(args, format);  // sets pointer to the beginning of our args
+
+  fprintf(stream, "[%s] [%s] ", time_str, get_log_level_string(level));
+  vfprintf(stream, format, args);
+  fprintf(stream, "\n");
+
+  va_end(args);  // cleans up va args memory
+}