use tuphium::activations::Activation;
use tuphium::layers::{Layer, Module};

fn run_one_step(module: &mut dyn Module, input: &Vec<f32>) -> Vec<f32> {
    module.zero_grad();
    let out = module.forward(input).unwrap();
    let ones = vec![1.0; out.len()];
    module.backward(&ones).unwrap();
    module.grads()
}

#[test]
fn module_trait_object_works() {
    let mut layer = Layer::new(3, 2, Activation::Sigmoid).unwrap();
    let input = vec![0.1, 0.2, 0.3];
    let grads = run_one_step(&mut layer, &input);
    assert_eq!(grads.len(), layer.params().len());
}

#[test]
fn zero_grad_via_trait_object_resets_grads() {
    let mut layer: Box<dyn Module> = Box::new(Layer::new(2, 2, Activation::ReLU).unwrap());
    let input = vec![1.0, 1.0];
    layer.forward(&input).unwrap();
    layer.backward(&vec![1.0, 1.0]).unwrap();
    layer.zero_grad();
    for g in layer.grads() {
        assert_eq!(g, 0.0);
    }
}
