[Fix] : Correct Softmax Gradient

include/cten.h

@@ -35,6 +35,7 @@ typedef struct GradNode {
     struct Tensor inputs[4];
     int n_inputs;
     const char* name;
+    int params[4];
 } GradNode;
 
 typedef struct {
@@ -111,7 +112,7 @@ Tensor nn_sigmoid(Tensor input);
 Tensor nn_tanh(Tensor input);
 Tensor nn_elu(Tensor self, float alpha);
 Tensor nn_selu(Tensor self);
-Tensor nn_softmax(Tensor input);
+Tensor nn_softmax(Tensor input, int dim);
 Tensor Glorot_init(TensorShape shape, bool requires_grad);
 Tensor nn_crossentropy(Tensor y_true, Tensor y_pred);
 Tensor nn_softmax_crossentropy(Tensor y_true, Tensor logits);

src/basic.c

@@ -175,7 +175,9 @@ void Tensor_backward(Tensor self, Tensor grad) {
 
         // Step 2: Apply the chain rule (upstream_grad * local_grad)
         Tensor combined_grad;
-        if(strcmp(self.node->name, "Matmul") == 0) {
+        if (strcmp(self.node->name, "Softmax") == 0) {
+            combined_grad = input_grad;
+        } else if(strcmp(self.node->name, "Matmul") == 0) {
             if (i == 0) {
                 combined_grad = Tensor_matmul(grad, input_grad);
             } else {

src/nn.c

@@ -302,57 +302,76 @@ static Tensor GradFn_softmax(Tensor self, int i) {
     Tensor input = self.node->inputs[i];
     Tensor grad = Tensor_new(input.shape, false);
 
-    int dim = TensorShape_dim(self.shape);
-    int batch_size = self.shape[0];
-    int num_classes = self.shape[1];  
-    for(int b = 0; b < batch_size; b++){
-        for(int i = 0; i < num_classes; i++) {
-            for(int j = 0; j < num_classes; j++) {
-                float softmax_i = self.data->flex[b * num_classes + i];
-                float softmax_j = self.data->flex[b * num_classes + j];
-                float value;
-                if(i == j){
-                    value = softmax_i * (1.0f - softmax_i);
-                } 
-                else{
-                    value = -softmax_i * softmax_j;
-                }
-
-                if(i == j){
-                    grad.data->flex[b * num_classes + i] = value;
-                }
+    int dim = self.node->params[0];
+    int input_ndim = TensorShape_dim(input.shape);
+
+    int dim_size = self.shape[dim];
+    int outer_size = 1;
+    for(int j = 0; j < dim; j++) {
+        outer_size *= self.shape[j];
+    }
+    int inner_size = 1;
+    for(int j = dim + 1; j < input_ndim; j++) {
+        inner_size *= self.shape[j];
+    }
+
+    float* s_data = self.data->flex; // Softmax output data (s)
+    float* upstream_grad_data = self.node->grad.data->flex; // Upstream grad (dL/ds)
+    float* input_grad_data = grad.data->flex; // Resulting grad (dL/dz)
+    for (int outer = 0; outer < outer_size; outer++) {
+        for (int inner = 0; inner < inner_size; inner++) {
+            int slice_offset = outer * dim_size * inner_size + inner;
+            // Step 1. Calculate the dot product for the current slice: sum_k(dL/ds_k * s_k)
+            float dot_product = 0.0f;
+            for (int k = 0; k < dim_size; k++) {
+                int index = slice_offset + k * inner_size;
+                dot_product += upstream_grad_data[index] * s_data[index];
+            }
+
+            // Step 2. Calculate the final gradient using the formula: dL/dz_j = s_j * (dL/ds_j - dot_product)
+            for (int k = 0; k < dim_size; k++) {
+                int index = slice_offset + k * inner_size;
+                input_grad_data[index] = s_data[index] * (upstream_grad_data[index] - dot_product);
             }
         }
     }
     return grad;
 }
 
-Tensor nn_softmax(Tensor self) {
+Tensor nn_softmax(Tensor self, int dim) {
     bool requires_grad = !cten_is_eval() && self.node != NULL;
     Tensor res = Tensor_new(self.shape, requires_grad);
     int self_dim = TensorShape_dim(self.shape);
-    assert(self_dim > 0);
-    int last_dim_size = self.shape[self_dim - 1];
-    int outer_size = self.data->numel / last_dim_size;
-
+    assert(dim >= 0 && dim < self_dim);
+    int dim_size = self.shape[dim];
+    int outer_size = 1;
+    for(int i = 0; i < dim; i++) {
+        outer_size *= self.shape[i];
+    }
+    int inner_size = 1;
+    for(int i = dim + 1; i < self_dim; i++) {
+        inner_size *= self.shape[i];
+    }
+
     for(int outer = 0; outer < outer_size; outer++) {
-        float max_val = -INFINITY;
-        float sum = 0;
-
-        for(int d = 0; d < last_dim_size; d++) {
-            int index = outer * last_dim_size + d;
-            max_val = fmaxf(max_val, self.data->flex[index]);
-        }
-
-        for(int d = 0; d < last_dim_size; d++) {
-            int index = outer * last_dim_size + d;
-            res.data->flex[index] = expf(self.data->flex[index] - max_val);
-            sum += res.data->flex[index];
-        }
-
-        for(int d = 0; d < last_dim_size; d++) {
-            int index = outer * last_dim_size + d;
-            res.data->flex[index] /= sum;
+        for(int inner = 0; inner < inner_size; inner++) {
+            int slice_offset = outer * dim_size * inner_size + inner;
+            float max_val = -INFINITY;
+            for(int k = 0; k < dim_size; k++) {
+                int index = slice_offset + k * inner_size;
+                max_val = fmaxf(max_val, self.data->flex[index]);
+            }
+            float sum = 0.0f;
+            for(int k = 0; k < dim_size; k++) {
+                int index = slice_offset + k * inner_size;
+                float val = expf(self.data->flex[index] - max_val);
+                res.data->flex[index] = val;
+                sum += val;
+            }
+            for(int k = 0; k < dim_size; k++) {
+                int index = slice_offset + k * inner_size;
+                res.data->flex[index] /= sum;
+            }
         }
     }
 
@@ -361,6 +380,7 @@ Tensor nn_softmax(Tensor self) {
         res.node->inputs[0] = self;
         res.node->n_inputs = 1; 
         res.node->name = "Softmax";     
+        res.node->params[0] = dim;
     }
     return res;
 }
@@ -482,8 +502,9 @@ static Tensor GradFn_softmax_crossentropy(Tensor self, int i) {
 Tensor nn_softmax_crossentropy(Tensor y_true, Tensor logits) {
     bool requires_grad = !cten_is_eval() && logits.node != NULL;
     //disable gradient computation
-    cten_begin_eval(); 
-    Tensor y_pred = nn_softmax(logits);
+    cten_begin_eval();
+    int last_dim_logits = TensorShape_dim(logits.shape) - 1;
+    Tensor y_pred = nn_softmax(logits, last_dim_logits);
     Tensor loss = nn_crossentropy(y_true, y_pred);
     cten_end_eval();
     Tensor res = Tensor_zeros((TensorShape){1}, requires_grad);

src/operator.c

@@ -82,7 +82,6 @@ Tensor Tensor_mul(Tensor self, Tensor other) {
     return res;
 }
 
-
 Tensor Tensor_mulf(Tensor self, float other) {
     Tensor tmp = Tensor_new(self.shape, false);
     for(int i = 0; i < tmp.data->numel; i++) {
@@ -283,7 +282,6 @@ static Tensor GradFn_sub(Tensor self, int i) {
     return res;
 }
 
-
 static Tensor GradFn_div(Tensor self, int i) {
     Tensor res = Tensor_new(self.shape, false);
     Tensor x = self.node->inputs[0];