Function adam 

pub fn adam(
    grad: fn(&[f64]) -> Vec<f64>,
    x0: &[f64],
    lr: f64,
    beta1: f64,
    beta2: f64,
    max_iter: usize,
    tol: f64,
) -> Vec<f64>