use scirs2_core::ndarray::{arr1, arr2};
use scirs2_neural::activations_minimal::{Activation, Softmax};
#[allow(dead_code)]
fn main() {
println!("Testing softmax implementation...\n");
let input = arr1(&[1.0, 2.0, 3.0]);
println!("Input: {input:?}");
let softmax = Softmax::new(0);
let output = softmax
.forward(&input.clone().into_dyn())
.expect("Operation failed");
println!("Softmax output: {output:?}");
let sum: f64 = output.sum();
println!("Sum of softmax: {sum}");
assert!((sum - 1.0).abs() < 1e-6, "Softmax should sum to 1");
println!("\nTest case 2: 2D batch");
let input_2d = arr2(&[[1.0, 2.0, 3.0], [3.0, 2.0, 1.0], [2.0, 2.0, 2.0]]);
println!("Input 2D:\n{input_2d:?}");
let softmax_2d = Softmax::new(1);
let output_2d = softmax_2d
.forward(&input_2d.clone().into_dyn())
.expect("Operation failed");
println!("Softmax output 2D:\n{output_2d:?}");
for i in 0..output_2d.shape()[0] {
let row_sum: f64 = output_2d.slice(scirs2_core::ndarray::s![i, ..]).sum();
println!("Row {i} sum: {row_sum}");
assert!((row_sum - 1.0).abs() < 1e-6, "Each row should sum to 1");
}
println!("\nTest case 3: Gradient computation");
let grad_output = arr1(&[0.1, 0.2, 0.3]).into_dyn();
let forward_output = softmax
.forward(&input.clone().into_dyn())
.expect("Operation failed");
let grad_input = softmax
.backward(&grad_output, &forward_output)
.expect("Operation failed");
println!("Gradient input: {grad_input:?}");
println!("\nAll tests passed!");
}