Function adam_step 

pub fn adam_step(
    embedding: &mut [[f32; 2]],
    grad: &[[f32; 2]],
    state: &mut AdamState,
    t: usize,
    lr: f32,
)

Apply one Adam step: update embedding in place.

t is the 1-indexed iteration number (used for bias correction).