Improve activation functions and tensors

include/nn_activation.h

@@ -35,13 +35,38 @@ typedef enum {
 } NNActFuncType;
 
 /**
- * @brief Represents the activation function.
+ * @brief Represents an activation function.
  */
-typedef union {
-    NNActFuncScalar scalar_func;
-    NNActFuncTensor tensor_func;
+typedef struct {
+    NNActFuncType type;
+    union {
+        NNActFuncScalar scalar_func;
+        NNActFuncTensor tensor_func;
+    };
 } NNActFunc;
 
+/**
+ * @brief Initializes a new activation function.
+ *
+ * @param type The type of activation function (scalar or tensor).
+ * @param func The activation function.
+ *
+ * @return The activation function.
+ */
+NNActFunc nn_act_func_init(NNActFuncType type, void *func);
+
+/**
+ * @brief Computes an activation function on the input tensor and stores the result in the output tensor.
+ *
+ * @param act_func The activation function.
+ * @param input The input tensor.
+ * @param output The output tensor.
+ * @param error The error instance to set if an error occurs.
+ *
+ * @return True or false.
+ */
+bool nn_act_func(NNActFunc act_func, const NNTensor *input, NNTensor *output, NNError *error);
+
 /**
  * @brief Returns the identity activation function result.
  *

include/nn_layer.h

@@ -50,7 +50,6 @@ typedef struct {
     NNActFunc act_func;
     NNTensor *weights;
     NNTensor *biases;
-
 } NNLayer;
 
 /**

include/nn_tensor.h

@@ -239,4 +239,16 @@ NN_TENSOR_DEFINE_DESTROY(NNTensor);
 #endif
 #endif // NN_TENSOR_DEFINED
 
+/**
+ * @brief Slice the input tensor and store the slice in the output tensor.
+ *
+ * @param input The input tensor to slice.
+ * @param offset The offset of the slice.
+ * @param sizes The sizes of the slice.
+ * @param output The tensor to store the slice.
+ *
+ * @return void
+ */
+void nn_tensor_slice(const NNTensor *input, const size_t offset, const size_t *sizes, NNTensor *output);
+
 #endif // NN_TENSOR_H

scripts/test/layer_gen_tc.py

@@ -36,11 +36,11 @@ def generate_test_case(batch_size, inputs_size, output_size, act_func_name):
 
     # Determine activation function
     if act_func_name == "softmax":
-        act_func_c = f".tensor_func = {act_func_map[act_func_name]}"
+        act_func_c = act_func_map[act_func_name]
         act_func_type = "NN_ACT_FUNC_TENSOR"
         act_func = nn_act_func_softmax
     else:
-        act_func_c = f".scalar_func = {act_func_map[act_func_name]}"
+        act_func_c = act_func_map[act_func_name]
         act_func_type = "NN_ACT_FUNC_SCALAR"
         act_func = globals()[f"nn_act_func_{act_func_name}"]
 
@@ -60,8 +60,7 @@ def generate_test_case(batch_size, inputs_size, output_size, act_func_name):
         .output_size = {output_size},
         .mat_mul_func = nn_mat_mul,
         .mat_transpose_func = nn_mat_transpose,
-        .act_func_type = {act_func_type},
-        .act_func = {{ {act_func_c} }},
+        .act_func = nn_act_func_init({act_func_type}, {act_func_c}),
         .weights = nn_tensor_init_NNTensor(2, (const size_t[]){{{output_size}, {inputs_size}}}, false, (const NNTensorUnit[]){{{weights_c}}}, NULL),
         .biases = nn_tensor_init_NNTensor(1, (const size_t[]){{{output_size}}}, false, (const NNTensorUnit[]){{{biases_c}}}, NULL),
         .inputs = nn_tensor_init_NNTensor(2, (const size_t[]){{{batch_size}, {inputs_size}}}, false, (const NNTensorUnit[]){{{inputs_c}}}, NULL),

scripts/test/layer_multi_gen_tc.py

@@ -87,10 +87,10 @@ def generate_test_cases(batch_sizes, input_sizes, output_sizes, hidden_layers):
                     layer_configs_c = []
                     for layer in layers:
                         if layer.act_func_name == "softmax":
-                            act_func_c = f".tensor_func = {act_func_map[layer.act_func_name]}"
+                            act_func_c = act_func_map[layer.act_func_name]
                             act_func_type = "NN_ACT_FUNC_TENSOR"
                         else:
-                            act_func_c = f".scalar_func = {act_func_map[layer.act_func_name]}"
+                            act_func_c = act_func_map[layer.act_func_name]
                             act_func_type = "NN_ACT_FUNC_SCALAR"
 
                         weights_c = ", ".join(map(str, layer.weights.flatten())) if layer.weights is not None else None
@@ -102,10 +102,9 @@ def generate_test_cases(batch_sizes, input_sizes, output_sizes, hidden_layers):
                             .layer_type = {layer.layer_type},
                             .inputs_size = {layer.inputs_size},
                             .output_size = {layer.output_size},
-                            .mat_mul_func = nn_mat_mul,
-                            .mat_transpose_func = nn_mat_transpose,
-                            .act_func_type = {act_func_type},
-                            .act_func = {{ {act_func_c} }},
+                            .mat_mul_func = {'NULL' if layer.layer_idx == 0 else 'nn_mat_mul'},
+                            .mat_transpose_func = {'NULL' if layer.layer_idx == 0 else 'nn_mat_transpose'},
+                            .act_func = nn_act_func_init({act_func_type}, {act_func_c}),
                             .weights = {"nn_tensor_init_NNTensor(2, (const size_t[]){" + str(layer.output_size) + ", " + str(layer.inputs_size) + "}, false, (const NNTensorUnit[]){ " + weights_c + "}, NULL)" if weights_c is not None else 'NULL'},
                             .biases = {"nn_tensor_init_NNTensor(1, (const size_t[]){" + str(layer.output_size) + "}, false, (const NNTensorUnit[]){ " + biases_c + "}, NULL)" if biases_c is not None else 'NULL'},
                             .inputs = {"nn_tensor_init_NNTensor(2, (const size_t[]){" + str(batch_size) + ", " + str(layer.inputs_size) + "}, false, (const NNTensorUnit[]){ " + inputs_c + "}, NULL)" if inputs_c is not None else 'NULL'},

src/nn_activation.c

@@ -8,6 +8,33 @@
 
 // TODO: Add tests
 
+NNActFunc nn_act_func_init(NNActFuncType type, void *func) {
+    NN_DEBUG_PRINT(5, "function %s called with type=%d\n", __func__, type);
+
+    NNActFunc act_func;
+    act_func.type = type;
+    if (type == NN_ACT_FUNC_SCALAR) {
+        act_func.scalar_func = (NNActFuncScalar)func;
+    } else if (type == NN_ACT_FUNC_TENSOR) {
+        act_func.tensor_func = (NNActFuncTensor)func;
+    }
+    return act_func;
+}
+
+bool nn_act_func(NNActFunc act_func, const NNTensor *input, NNTensor *output, NNError *error) {
+    NN_DEBUG_PRINT(5, "function %s called with input.dims=%zu output.dims=%zu\n", __func__, input->dims, output->dims);
+
+    switch (act_func.type) {
+    case NN_ACT_FUNC_SCALAR:
+        return nn_act_func_scalar_batch(act_func.scalar_func, input, output, error);
+    case NN_ACT_FUNC_TENSOR:
+        return nn_act_func_tensor_batch(act_func.tensor_func, input, output, error);
+    default:
+        nn_error_set(error, NN_ERROR_INVALID_ARGUMENT, "invalid activation function type");
+        return false;
+    }
+}
+
 NNTensorUnit nn_act_func_identity(NNTensorUnit n) {
     NN_DEBUG_PRINT(5, "function %s called with n=%f\n", __func__, n);
 
@@ -68,13 +95,6 @@ bool nn_act_func_softmax(const NNTensor *input, NNTensor *output, NNError *error
     return true;
 }
 
-static void nn_tensor_create_slice(const NNTensor *tensor, const size_t offset, const size_t *sizes, NNTensor *slice) {
-    slice->flags = NN_TENSOR_FLAG_INIT;
-    slice->dims = 1;
-    slice->sizes = (size_t *)sizes;
-    slice->data = &tensor->data[offset];
-}
-
 bool nn_act_func_scalar_batch(const NNActFuncScalar act_func, const NNTensor *input, NNTensor *output, NNError *error) {
     NN_DEBUG_PRINT(5, "function %s called with input.dims=%zu output.dims=%zu\n", __func__, input->dims, output->dims);
 
@@ -92,7 +112,7 @@ bool nn_act_func_scalar_batch(const NNActFuncScalar act_func, const NNTensor *in
     size_t sizes[1] = {sample_size};
     NNTensor input_slice;
     for (size_t i = 0; i < batch_size; i++) {
-        nn_tensor_create_slice(input, i * sample_size, sizes, &input_slice);
+        nn_tensor_slice(input, i * sample_size, sizes, &input_slice);
         for (size_t j = 0; j < sample_size; ++j) {
             output->data[i * sample_size + j] = act_func(input_slice.data[j]);
         }
@@ -120,8 +140,8 @@ bool nn_act_func_tensor_batch(const NNActFuncTensor act_func, const NNTensor *in
     NNTensor output_slice;
 
     for (size_t i = 0; i < batch_size; i++) {
-        nn_tensor_create_slice(input, i * sample_size, sizes, &input_slice);
-        nn_tensor_create_slice(output, i * sample_size, sizes, &output_slice);
+        nn_tensor_slice(input, i * sample_size, sizes, &input_slice);
+        nn_tensor_slice(output, i * sample_size, sizes, &output_slice);
 
         if (!act_func(&input_slice, &output_slice, error)) {
             return false;

src/nn_layer.c

@@ -61,6 +61,7 @@ NNLayer *nn_layer_init(size_t input_size, size_t output_size, NNError *error) {
         return NULL;
     }
     layer->flags = NN_LAYER_FLAG_INIT;
+    nn_layer_check_forward_ready(layer); // update flags
 
     return layer;
 }
@@ -158,8 +159,9 @@ bool nn_layer_forward(const NNLayer *layer, const NNTensor *inputs, NNTensor *ou
         return false;
     }
 
-    // Check if matrix multiplication function is set
+    // Check if weights are set
     if (layer->flags & NN_LAYER_FLAG_WEIGHTS_SET) {
+        // Check if matrix multiplication function is set
         if (layer->flags & NN_LAYER_FLAG_MAT_MUL_FUNC_SET) {
             // Check if matrix transpose function is set
             if (layer->flags & NN_LAYER_FLAG_MAT_TRANSPOSE_FUNC_SET) {
@@ -168,6 +170,7 @@ bool nn_layer_forward(const NNLayer *layer, const NNTensor *inputs, NNTensor *ou
                 if (!weights) {
                     return false;
                 }
+                // TODO: Should we overwrite the weights tensor so that we don't have to allocate a new tensor every time?
                 if (!nn_mat_transpose(layer->weights, weights, error)) {
                     return false;
                 }

src/nn_mat_mul.c

@@ -11,6 +11,7 @@ bool nn_mat_mul(const NNTensor *a, const NNTensor *b, NNTensor *result, NNError
         nn_error_set(error, NN_ERROR_INVALID_ARGUMENT, "tensor a or b is not initialized");
         return NULL;
     } else if (a->dims != 2 || b->dims != 2 || a->sizes[1] != b->sizes[0]) {
+        // Number of columns in the first tensor must be equal to the number of rows in the second tensor
         nn_error_setf(error, NN_ERROR_INVALID_ARGUMENT, "only 2-dimensional tensors with matching inner dimensions are allowed, a.dims=%zu b.dims=%zu a.sizes[1]=%zu b.sizes[0]=%zu", a->dims, b->dims, a->sizes[1], b->sizes[0]);
         return NULL;
     }

src/nn_tensor.c

@@ -0,0 +1,8 @@
+#include "nn_tensor.h"
+
+void nn_tensor_slice(const NNTensor *input, const size_t offset, const size_t *sizes, NNTensor *output) {
+    output->flags = NN_TENSOR_FLAG_INIT;
+    output->dims = 1;
+    output->sizes = (size_t *)sizes;
+    output->data = &input->data[offset];
+}

tests/arch/generic/layer/include.txt

@@ -1,8 +1,9 @@
 src/nn_activation.c
 src/nn_app.c
 src/nn_config.c
-src/nn_mat_mul.c
-src/nn_mat_transpose.c
 src/nn_error.c
 src/nn_layer.c
+src/nn_mat_mul.c
+src/nn_mat_transpose.c
+src/nn_tensor.c
 src/nn_test.c