Function adam_state_to_map 

pub fn adam_state_to_map(
    state: &[(u64, Tensor, Tensor, u64)],
) -> HashMap<String, Tensor>

Flatten Adam/AdamW state into a string-keyed map for serialization.

Keys: "adam.{param_id}.m", "adam.{param_id}.v", "adam.{param_id}.step"