Fixing CPU Tensor Softmax - Needs Testing

hpvm/projects/hpvm-tensor-rt/tensor_runtime/src/tensor_cpu_runtime.cc

@@ -1026,32 +1026,34 @@ void *tensorSoftmaxCPU(void *input_ptr) {
   int n = input->dims.dim_sizes[0];
   int c = input->dims.dim_sizes[1];
 
-  float max = logits[0];
-  for (unsigned int i = 0; i < n * c; i++){
-    if (logits[i] > max){
-      max = logits[i];
-    }
-  }
   
   omp_set_num_threads(4);
 #pragma omp parallel for
   for (int i = 0; i < n; i++) {
+
+    float max = logits[i * c];
+    for (unsigned int k = i * c; k < c + i * c; k++){
+      if (logits[k] > max){
+        max = logits[k];
+      }
+    }
+  
     double x = 0;
-    for (int j = i * c; j < c + i * c; j++) {
-      logits[j] = exp(logits[j] / max );
+    for (int j = i * c; j < c + i * c; j++) {   
+      logits[j] = exp( logits[j] - max );
     }
 
 #pragma omp simd reduction(+ : x)
     for (int j = i * c; j < i * c + c; j++) {
       x += logits[j];
     }
-
-    //printf("x = %f \n ", x);
     
 #pragma omp simd
     for (int j = i * c; j < i * c + c; j++) {
       logits[j] /= x;
     }
+
+    //printf("logits[i * c] = %f \n ", logits[i * c]);
   }
   
   return input;