fix NaN/Inf issue with exp

Author: mseminatore

ann.c

@@ -278,25 +278,29 @@ static void record_history(PNetwork pnet, real loss, real learning_rate);
 //------------------------------
 static void softmax(PNetwork pnet)
 {
-	real one_over_sum = 0.0;
 	real sum = 0.0;
 
 	// find the sum of the output node values, excluding the bias node
 	int output_layer = pnet->layer_count - 1;
 	int node_count = pnet->layers[output_layer].node_count;
-	
-	//tensor_exp(pnet->layers[output_layer].t_values);
-	//one_over_sum = (real)1.0 / (tensor_sum(pnet->layers[output_layer].t_values) - 1.0);
-	//tensor_mul_scalar(pnet->layers[output_layer].t_values, one_over_sum);
+	real *values = pnet->layers[output_layer].t_values->values;
+
+	// Find max for numerical stability
+	real max_val = values[0];
+	for (int node = 1; node < node_count; node++)
+	{
+		if (values[node] > max_val)
+			max_val = values[node];
+	}
 
 	for (int node = 0; node < node_count; node++)
 	{
-		sum += (real)exp(pnet->layers[output_layer].t_values->values[node]);
+		sum += (real)exp(values[node] - max_val);
 	}
 
 	for (int node = 0; node < node_count; node++)
 	{
-		pnet->layers[output_layer].t_values->values[node] = (real)(exp(pnet->layers[output_layer].t_values->values[node]) / sum);
+		values[node] = (real)(exp(values[node] - max_val) / sum);
 	}
 }
 
@@ -543,7 +547,7 @@ static real compute_cross_entropy(PNetwork pnet, PTensor outputs)
 
 	for (int i = 0; i < poutput_layer->node_count; i++)
 	{
-		xe += (real)(outputs->values[i] * log(poutput_layer->t_values->values[i]));
+		xe += (real)(outputs->values[i] * log(fmax(poutput_layer->t_values->values[i], (real)1e-7)));
 	}
 
 	return -xe;
@@ -598,10 +602,18 @@ static void softmax_batched(PNetwork pnet, int batch_size)
 		real sum = (real)0.0;
 		int row_offset = b * node_count;
 
-		// Compute exp and sum for this sample
+		// Find max for numerical stability
+		real max_val = batch->values[row_offset];
+		for (int n = 1; n < node_count; n++)
+		{
+			if (batch->values[row_offset + n] > max_val)
+				max_val = batch->values[row_offset + n];
+		}
+
+		// Compute exp(x - max) and sum for this sample
 		for (int n = 0; n < node_count; n++)
 		{
-			batch->values[row_offset + n] = (real)exp(batch->values[row_offset + n]);
+			batch->values[row_offset + n] = (real)exp(batch->values[row_offset + n] - max_val);
 			sum += batch->values[row_offset + n];
 		}
 
@@ -3289,7 +3301,11 @@ static void write_json_float_array(FILE *fptr, const real *values, int count)
 	for (int i = 0; i < count; i++)
 	{
 		if (i > 0) fprintf(fptr, ", ");
-		fprintf(fptr, "%.8g", (double)values[i]);
+		double v = (double)values[i];
+		if (isnan(v) || isinf(v))
+			fprintf(fptr, "0.0");
+		else
+			fprintf(fptr, "%.8g", v);
 	}
 	fprintf(fptr, "]");
 }

test_onnx_export.c

@@ -288,6 +288,30 @@ void test_main(int argc, char *argv[]) {
 
     remove_file(import_file);
 
+    // ========================================================================
+    // NaN/Inf HANDLING TESTS
+    // ========================================================================
+    SUITE("ONNX Export NaN/Inf Handling");
+    COMMENT("Testing that NaN/Inf values are serialized as valid JSON...");
+
+    net = ann_make_network(OPT_SGD, LOSS_MSE);
+    ann_add_layer(net, 2, LAYER_INPUT, ACTIVATION_NULL);
+    ann_add_layer(net, 2, LAYER_OUTPUT, ACTIVATION_SIGMOID);
+
+    // Inject NaN and Inf into weights
+    net->layers[0].t_weights->values[0] = NAN;
+    net->layers[0].t_weights->values[1] = INFINITY;
+    net->layers[0].t_bias->values[0] = -INFINITY;
+
+    result = ann_export_onnx(net, test_file);
+    TESTEX("Export with NaN/Inf values returns ERR_OK", (result == ERR_OK));
+    TESTEX("Output does not contain 'nan'", !file_contains(test_file, "nan"));
+    TESTEX("Output does not contain 'inf'", !file_contains(test_file, "inf"));
+    TESTEX("Output contains 0.0 replacement", file_contains(test_file, "0.0"));
+
+    remove_file(test_file);
+    ann_free_network(net);
+
     // Cleanup any remaining test files
     remove_file(test_file);
     remove_file(pikchr_file);