pub fn normal_log_prob(x: f32, mean: f32, std: f32) -> f32
Compute log_prob for a normal distribution.