Add softmax and refactor activation functions

examples/arch/generic/layer/main.c

@@ -16,7 +16,7 @@ int main() {
     const int batch_size = 2;
 
     // Initialize a layer with the given input and output sizes, ReLU activation function, and dot product function
-    if (!nn_layer_init(&layer, input_size, output_size, nn_activation_func_relu, nn_dot_product, &error)) {
+    if (!nn_layer_init(&layer, input_size, output_size, &error)) {
         fprintf(stderr, "error: %s\n", error.message);
         return 1;
     }
@@ -33,6 +33,19 @@ int main() {
         return 1;
     }
 
+    // Set the dot product function of the layer
+    if (!nn_layer_set_dot_product_func(&layer, nn_dot_product, &error)) {
+        fprintf(stderr, "error: %s\n", error.message);
+        return 1;
+    }
+
+    // Set the activation function of the layer
+    NNActivationFunction act_func = {.scalar = nn_activation_func_relu};
+    if (!nn_layer_set_activation_func(&layer, act_func, &error)) {
+        fprintf(stderr, "error: %s\n", error.message);
+        return 1;
+    }
+
     // Generate random inputs
     float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE];
     for (size_t i = 0; i < batch_size; ++i) {

include/nn_activation.h

@@ -1,8 +1,21 @@
-#ifndef NN_ACTIVATION_FUNCTIONS_H
-#define NN_ACTIVATION_FUNCTIONS_H
+#ifndef NN_ACTIVATION_FUNCTION_H
+#define NN_ACTIVATION_FUNCTION_H
+
+#include "nn_error.h"
+#include <stdbool.h>
+#include <stddef.h>
+
+#ifndef NN_SOFTMAX_MAX_SIZE
+#define NN_SOFTMAX_MAX_SIZE 64
+#endif
 
 // NNActivationFunction represents an activation function.
-typedef float (*NNActivationFunction)(float);
+typedef float (*NNActivationFunctionScalar)(float);
+typedef bool (*NNActivationFunctionVector)(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error);
+typedef union {
+    NNActivationFunctionScalar scalar;
+    NNActivationFunctionVector vector;
+} NNActivationFunction;
 
 // nn_activation_func_identity returns x.
 float nn_activation_func_identity(float x);
@@ -13,4 +26,7 @@ float nn_activation_func_sigmoid(float x);
 // nn_activation_func_relu returns the ReLU of x.
 float nn_activation_func_relu(float x);
 
-#endif // NN_ACTIVATION_FUNCTIONS_H
+// nn_activation_func_softmax calculates the softmax of the input and stores the result in the output.
+bool nn_activation_func_softmax(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error);
+
+#endif // NN_ACTIVATION_FUNCTION_H

include/nn_error.h

@@ -6,8 +6,11 @@
 // NNErrorCode defines the error codes.
 typedef enum {
     NN_ERROR_NONE = 0,                // no error
+    NN_ERROR_NOT_IMPLEMENTED,         // not implemented
     NN_ERROR_INVALID_INSTANCE,        // invalid instance
     NN_ERROR_INVALID_SIZE,            // invalid size
+    NN_ERROR_INVALID_VALUE,           // invalid value
+    NN_ERROR_INVALID_TYPE,            // invalid type
     NN_ERROR_NEON_NOT_AVAILABLE,      // NEON instructions not available
     NN_ERROR_CMSIS_DSP_NOT_AVAILABLE, // CMSIS-DSP functions not available
 } NNErrorCode;

include/nn_layer.h

@@ -33,12 +33,12 @@ typedef struct {
     size_t output_size;
     float weights[NN_LAYER_MAX_OUTPUT_SIZE][NN_LAYER_MAX_INPUT_SIZE];
     float biases[NN_LAYER_MAX_BIASES];
-    NNActivationFunction act_func;
     NNDotProductFunction dot_product_func;
+    NNActivationFunction act_func;
 } NNLayer;
 
 // nn_layer_init initializes a layer with the given arguments.
-bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error);
+bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNError *error);
 
 // nn_layer_init_weights_gaussian initializes the weights of the layer with a Gaussian distribution.
 bool nn_layer_init_weights_gaussian(NNLayer *layer, float scale, NNError *error);
@@ -52,6 +52,12 @@ bool nn_layer_set_weights(NNLayer *layer, const float weights[NN_LAYER_MAX_OUTPU
 // nn_layer_set_biases sets the biases of the given layer.
 bool nn_layer_set_biases(NNLayer *layer, const float biases[NN_LAYER_MAX_BIASES], NNError *error);
 
+// nn_layer_set_dot_product_func sets the dot product function of the given layer.
+bool nn_layer_set_dot_product_func(NNLayer *layer, NNDotProductFunction dot_product_func, NNError *error);
+
+// nn_layer_set_activation_func sets the activation function of the given layer.
+bool nn_layer_set_activation_func(NNLayer *layer, NNActivationFunction act_func, NNError *error);
+
 // nn_layer_forward computes the given layer with the given inputs and stores the result in outputs.
 bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE], float outputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_OUTPUT_SIZE], size_t batch_size, NNError *error);
 

include/nn_neuron.h

@@ -19,19 +19,25 @@ typedef struct {
     float weights[NN_NEURON_MAX_WEIGHTS];
     size_t input_size;
     float bias;
-    NNActivationFunction act_func;
     NNDotProductFunction dot_product_func;
+    NNActivationFunctionScalar act_func;
 } NNNeuron;
 
 // nn_neuron_init initializes a neuron with the given arguments.
-bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error);
+bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNError *error);
 
 // nn_neuron_set_weights sets the weights of the given neuron.
 bool nn_neuron_set_weights(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], NNError *error);
 
 // nn_neuron_set_bias sets the bias of the given neuron.
 bool nn_neuron_set_bias(NNNeuron *neuron, float bias, NNError *error);
 
+// nn_neuron_set_dot_product_func sets the dot product function of the given neuron.
+bool nn_neuron_set_dot_product_func(NNNeuron *neuron, NNDotProductFunction dot_product_func, NNError *error);
+
+// nn_neuron_set_activation_func sets the activation function of the given neuron.
+bool nn_neuron_set_activation_func(NNNeuron *neuron, NNActivationFunctionScalar act_func, NNError *error);
+
 // nn_neuron_compute computes the given neuron and returns the output.
 float nn_neuron_compute(const NNNeuron *neuron, const float inputs[NN_NEURON_MAX_WEIGHTS], NNError *error);
 

src/nn_activation.c

@@ -1,8 +1,9 @@
 #include "nn_activation.h"
+#include "nn_error.h"
 #include <math.h>
+#include <stdbool.h>
 
 // TODO: Add tests
-// TODO: Add softmax activation function.
 
 // nn_activation_func_identity returns x.
 float nn_activation_func_identity(float x) {
@@ -18,3 +19,45 @@ float nn_activation_func_sigmoid(float x) {
 float nn_activation_func_relu(float x) {
     return fmaxf(0, x);
 }
+
+// nn_activation_func_softmax calculates the softmax of the input and stores the result in the output.
+bool nn_activation_func_softmax(const float input[NN_SOFTMAX_MAX_SIZE], float output[NN_SOFTMAX_MAX_SIZE], size_t input_size, NNError *error) {
+    nn_error_set(error, NN_ERROR_NONE, NULL);
+    if (input == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "input is NULL");
+        return false;
+    } else if (output == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "output is NULL");
+        return false;
+    } else if (input_size == 0 || input_size > NN_SOFTMAX_MAX_SIZE) {
+        nn_error_set(error, NN_ERROR_INVALID_SIZE, "invalid input size");
+        return false;
+    }
+
+    // Find the maximum input value
+    float max_input = input[0];
+    for (size_t i = 1; i < input_size; ++i) {
+        if (input[i] > max_input) {
+            max_input = input[i];
+        }
+    }
+
+    // Compute exp(input[i] - max_input) to prevent overflow
+    float sum = 0.0f;
+    for (size_t i = 0; i < input_size; ++i) {
+        output[i] = expf(input[i] - max_input);
+        sum += output[i];
+    }
+
+    if (sum == 0.0f) {
+        nn_error_set(error, NN_ERROR_INVALID_VALUE, "sum is zero");
+        return false;
+    }
+
+    // Normalize to form a probability distribution
+    for (size_t i = 0; i < input_size; ++i) {
+        output[i] /= sum;
+    }
+
+    return true;
+}

src/nn_layer.c

@@ -11,7 +11,7 @@
 #endif
 
 // nn_layer_init initializes a layer with the given arguments.
-bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error) {
+bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNError *error) {
     nn_error_set(error, NN_ERROR_NONE, NULL);
     if (layer == NULL) {
         nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
@@ -27,12 +27,6 @@ bool nn_layer_init(NNLayer *layer, size_t input_size, size_t output_size, NNActi
     }
     layer->input_size = input_size;
     layer->output_size = output_size;
-    if (act_func) {
-        layer->act_func = act_func;
-    }
-    if (dot_product_func) {
-        layer->dot_product_func = dot_product_func;
-    }
 
     return true;
 }
@@ -104,6 +98,30 @@ bool nn_layer_set_biases(NNLayer *layer, const float biases[NN_LAYER_MAX_BIASES]
     return true;
 }
 
+// nn_layer_set_dot_product_func sets the dot product function of the given layer.
+bool nn_layer_set_dot_product_func(NNLayer *layer, NNDotProductFunction dot_product_func, NNError *error) {
+    nn_error_set(error, NN_ERROR_NONE, NULL);
+    if (layer == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
+        return false;
+    }
+    layer->dot_product_func = dot_product_func;
+
+    return true;
+}
+
+// nn_layer_set_activation_func sets the activation function of the given layer.
+bool nn_layer_set_activation_func(NNLayer *layer, NNActivationFunction act_func, NNError *error) {
+    nn_error_set(error, NN_ERROR_NONE, NULL);
+    if (layer == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "layer is NULL");
+        return false;
+    }
+    layer->act_func = act_func;
+
+    return true;
+}
+
 // nn_layer_forward computes the given layer with the given inputs and stores the result in outputs.
 bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_INPUT_SIZE], float outputs[NN_LAYER_MAX_BATCH_SIZE][NN_LAYER_MAX_OUTPUT_SIZE], size_t batch_size, NNError *error) {
     nn_error_set(error, NN_ERROR_NONE, NULL);
@@ -123,8 +141,8 @@ bool nn_layer_forward(const NNLayer *layer, const float inputs[NN_LAYER_MAX_BATC
             if (layer->dot_product_func != NULL) {
                 outputs[i][j] += layer->dot_product_func(inputs[i], layer->weights[j], layer->input_size);
             }
-            if (layer->act_func != NULL) {
-                outputs[i][j] = layer->act_func(outputs[i][j]);
+            if (layer->act_func.scalar != NULL) {
+                outputs[i][j] = layer->act_func.scalar(outputs[i][j]);
             }
         }
     }

src/nn_neuron.c

@@ -7,7 +7,7 @@
 #include <stdio.h>
 
 // nn_neuron_init initializes a neuron with the given arguments.
-bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNActivationFunction act_func, NNDotProductFunction dot_product_func, NNError *error) {
+bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS], size_t input_size, float bias, NNError *error) {
     nn_error_set(error, NN_ERROR_NONE, NULL);
     if (neuron == NULL) {
         nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
@@ -21,12 +21,6 @@ bool nn_neuron_init(NNNeuron *neuron, const float weights[NN_NEURON_MAX_WEIGHTS]
         neuron->weights[i] = weights[i];
     }
     neuron->bias = bias;
-    if (act_func != NULL) {
-        neuron->act_func = act_func;
-    }
-    if (dot_product_func != NULL) {
-        neuron->dot_product_func = dot_product_func;
-    }
     return true;
 }
 
@@ -54,6 +48,28 @@ bool nn_neuron_set_bias(NNNeuron *neuron, float bias, NNError *error) {
     return true;
 }
 
+// nn_neuron_set_dot_product_func sets the dot product function of the given neuron.
+bool nn_neuron_set_dot_product_func(NNNeuron *neuron, NNDotProductFunction dot_product_func, NNError *error) {
+    nn_error_set(error, NN_ERROR_NONE, NULL);
+    if (neuron == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
+        return false;
+    }
+    neuron->dot_product_func = dot_product_func;
+    return true;
+}
+
+// nn_neuron_set_activation_func sets the activation function of the given neuron.
+bool nn_neuron_set_activation_func(NNNeuron *neuron, NNActivationFunctionScalar act_func, NNError *error) {
+    nn_error_set(error, NN_ERROR_NONE, NULL);
+    if (neuron == NULL) {
+        nn_error_set(error, NN_ERROR_INVALID_INSTANCE, "neuron is NULL");
+        return false;
+    }
+    neuron->act_func = act_func;
+    return true;
+}
+
 // nn_neuron_compute computes the given neuron and returns the output.
 float nn_neuron_compute(const NNNeuron *neuron, const float inputs[NN_NEURON_MAX_WEIGHTS], NNError *error) {
     nn_error_set(error, NN_ERROR_NONE, NULL);

tests/arch/arm/cmsis-dsp/neuron/main.c

@@ -29,7 +29,11 @@ void run_test_cases(TestCase *test_cases, int n_cases, char *info, NNDotProductF
         NNNeuron neuron;
         NNError error;
 
-        nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, nn_activation_func_identity, dot_product_func, &error);
+        nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, &error);
+        assert(error.code == NN_ERROR_NONE);
+        nn_neuron_set_dot_product_func(&neuron, dot_product_func, &error);
+        assert(error.code == NN_ERROR_NONE);
+        nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error);
         assert(error.code == NN_ERROR_NONE);
         const float output = nn_neuron_compute(&neuron, tc.inputs, &error);
         assert(error.code == NN_ERROR_NONE);

tests/arch/arm/cmsis-dsp/neuron_perf/main.c

@@ -32,7 +32,13 @@ int main(int argc, char *argv[]) {
         inputs[i] = (float)rand() / (float)RAND_MAX;
     }
 
-    if (!nn_neuron_init(&neuron, weights, input_size, bias, nn_activation_func_identity, nn_dot_product_cmsis, &error)) {
+    if (!nn_neuron_init(&neuron, weights, input_size, bias, &error)) {
+        printf("error: %s\n", error.message);
+        return 1;
+    } else if (!nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error)) {
+        printf("error: %s\n", error.message);
+        return 1;
+    } else if (!nn_neuron_set_dot_product_func(&neuron, nn_dot_product_cmsis, &error)) {
         printf("error: %s\n", error.message);
         return 1;
     }

tests/arch/arm/neon/neuron/main.c

@@ -29,7 +29,11 @@ void run_test_cases(TestCase *test_cases, int n_cases, char *info, NNDotProductF
         NNNeuron neuron;
         NNError error;
 
-        nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, nn_activation_func_identity, dot_product_func, &error);
+        nn_neuron_init(&neuron, tc.weights, tc.input_size, tc.bias, &error);
+        assert(error.code == NN_ERROR_NONE);
+        nn_neuron_set_dot_product_func(&neuron, dot_product_func, &error);
+        assert(error.code == NN_ERROR_NONE);
+        nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error);
         assert(error.code == NN_ERROR_NONE);
         const float output = nn_neuron_compute(&neuron, tc.inputs, &error);
         assert(error.code == NN_ERROR_NONE);

tests/arch/arm/neon/neuron_perf/main.c

@@ -33,7 +33,13 @@ int main(int argc, char *argv[]) {
         inputs[i] = (float)rand() / (float)RAND_MAX;
     }
 
-    if (!nn_neuron_init(&neuron, weights, input_size, bias, nn_activation_func_identity, nn_dot_product_neon, &error)) {
+    if (!nn_neuron_init(&neuron, weights, input_size, bias, &error)) {
+        printf("error: %s\n", error.message);
+        return 1;
+    } else if (!nn_neuron_set_activation_func(&neuron, nn_activation_func_identity, &error)) {
+        printf("error: %s\n", error.message);
+        return 1;
+    } else if (!nn_neuron_set_dot_product_func(&neuron, nn_dot_product_neon, &error)) {
         printf("error: %s\n", error.message);
         return 1;
     }