From 3cd8adc3dec0999cb837c37649436b28acb453c7 Mon Sep 17 00:00:00 2001 From: Advaitgaur004 Date: Fri, 25 Jul 2025 23:47:23 +0530 Subject: [PATCH 1/4] fix(nn): Correct softmax gradient calculation and backward pass --- src/basic.c | 4 +++- src/nn.c | 37 ++++++++++++++++++------------------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/src/basic.c b/src/basic.c index 15f6bdd..04f7a76 100644 --- a/src/basic.c +++ b/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 { diff --git a/src/nn.c b/src/nn.c index e7d6a16..0347708 100644 --- a/src/nn.c +++ b/src/nn.c @@ -303,25 +303,24 @@ static Tensor GradFn_softmax(Tensor self, int 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; - } - } + assert(dim > 0); + int last_dim_size = self.shape[dim - 1]; + int outer_size = self.data->numel / last_dim_size; + + 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++) { + int offset = outer * last_dim_size; + // 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 < last_dim_size; k++) { + dot_product += upstream_grad_data[offset + k] * s_data[offset + k]; + } + // Step 2. Calculate the final gradient for each element in the slice + for (int j = 0; j < last_dim_size; j++) { + int index = offset + j; + input_grad_data[index] = s_data[index] * (upstream_grad_data[index] - dot_product); } } return grad; From 9313168276caaf755cd0cff803415534498e6d76 Mon Sep 17 00:00:00 2001 From: Advaitgaur004 Date: Fri, 25 Jul 2025 23:53:56 +0530 Subject: [PATCH 2/4] ultra minor cleanup in operator.c - This is no fix afterall...but yeah it looks good now, dont judge my career on the basis of this commit later on :( --- src/operator.c | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/operator.c b/src/operator.c index c457b69..63c4f88 100644 --- a/src/operator.c +++ b/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]; From 6ea516c613f9cc5162696e7d89aef3fa4b663f4f Mon Sep 17 00:00:00 2001 From: Advaitgaur004 Date: Fri, 1 Aug 2025 13:16:13 +0530 Subject: [PATCH 3/4] extented the softmax to include dim as parameter --- src/nn.c | 98 ++++++++++++++++++++++++++++++++++---------------------- 1 file changed, 60 insertions(+), 38 deletions(-) diff --git a/src/nn.c b/src/nn.c index 0347708..a341c79 100644 --- a/src/nn.c +++ b/src/nn.c @@ -302,56 +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); - assert(dim > 0); - int last_dim_size = self.shape[dim - 1]; - int outer_size = self.data->numel / last_dim_size; + 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++) { - int offset = outer * last_dim_size; - // 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 < last_dim_size; k++) { - dot_product += upstream_grad_data[offset + k] * s_data[offset + k]; - } - // Step 2. Calculate the final gradient for each element in the slice - for (int j = 0; j < last_dim_size; j++) { - int index = offset + j; - input_grad_data[index] = s_data[index] * (upstream_grad_data[index] - dot_product); + 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; + } } } @@ -360,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; } @@ -481,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); From cff22a330aa05860044e70440054d5e09fb4e6ab Mon Sep 17 00:00:00 2001 From: Advaitgaur004 Date: Fri, 1 Aug 2025 13:19:52 +0530 Subject: [PATCH 4/4] fix definitions - nn_softmax and changed GradNode structure --- include/cten.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/include/cten.h b/include/cten.h index f45c432..e876ed0 100644 --- a/include/cten.h +++ b/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);