rand-half 0.1.0

// src/lib.rs
use half::{bf16, f16};
use rand::distributions::{Distribution, Standard};
use rand::prelude::*;
use rand_distr::StandardNormal;
use std::fmt;

/// Wrapper type for bf16 to allow implementing foreign traits
#[derive(Debug, Clone, Copy)]
pub struct Bf16Wrapper(pub bf16);

/// Wrapper type for f16 to allow implementing foreign traits
#[derive(Debug, Clone, Copy)]
pub struct F16Wrapper(pub f16);

// Implement conversion methods
impl From<bf16> for Bf16Wrapper {
    fn from(value: bf16) -> Self {
        Bf16Wrapper(value)
    }
}

impl From<Bf16Wrapper> for bf16 {
    fn from(wrapper: Bf16Wrapper) -> Self {
        wrapper.0
    }
}

impl From<f16> for F16Wrapper {
    fn from(value: f16) -> Self {
        F16Wrapper(value)
    }
}

impl From<F16Wrapper> for f16 {
    fn from(wrapper: F16Wrapper) -> Self {
        wrapper.0
    }
}

// Implement display for convenience
impl fmt::Display for Bf16Wrapper {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", f32::from(self.0))
    }
}

impl fmt::Display for F16Wrapper {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", f32::from(self.0))
    }
}

// Now implement Distribution for our wrapper types
impl Distribution<Bf16Wrapper> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Bf16Wrapper {
        // Generate an f32 and convert to bf16
        let val: f32 = rng.gen();
        Bf16Wrapper(bf16::from_f32(val))
    }
}

impl Distribution<F16Wrapper> for Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F16Wrapper {
        // Generate an f32 and convert to f16
        let val: f32 = rng.gen();
        F16Wrapper(f16::from_f32(val))
    }
}

impl Distribution<Bf16Wrapper> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Bf16Wrapper {
        let x: f32 = StandardNormal.sample(rng);
        Bf16Wrapper(bf16::from_f32(x))
    }
}

impl Distribution<F16Wrapper> for StandardNormal {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> F16Wrapper {
        let x: f32 = StandardNormal.sample(rng);
        F16Wrapper(f16::from_f32(x))
    }
}

// Helper functions for generating uniform distributions
pub fn rand_uniform_bf16<R: Rng + ?Sized>(rng: &mut R, min: f32, max: f32) -> bf16 {
    let range = rand::distributions::Uniform::new(min, max);
    let val: f32 = range.sample(rng);
    bf16::from_f32(val)
}

pub fn rand_uniform_f16<R: Rng + ?Sized>(rng: &mut R, min: f32, max: f32) -> f16 {
    let range = rand::distributions::Uniform::new(min, max);
    let val: f32 = range.sample(rng);
    f16::from_f32(val)
}

pub fn rand_normal_bf16<R: Rng + ?Sized>(rng: &mut R, mean: f32, std: f32) -> bf16 {
    let normal = rand_distr::Normal::new(mean, std).unwrap();
    let val: f32 = normal.sample(rng);
    bf16::from_f32(val)
}

pub fn rand_normal_f16<R: Rng + ?Sized>(rng: &mut R, mean: f32, max: f32) -> f16 {
    let normal = rand_distr::Normal::new(mean, max).unwrap();
    let val: f32 = normal.sample(rng);
    f16::from_f32(val)
}

// Extension trait for RngCore to add convenience methods
pub trait HalfRngExt {
    fn gen_bf16(&mut self) -> bf16;
    fn gen_f16(&mut self) -> f16;
    fn gen_range_bf16(&mut self, min: f32, max: f32) -> bf16;
    fn gen_range_f16(&mut self, min: f32, max: f32) -> f16;
}

impl<R: RngCore + ?Sized> HalfRngExt for R {
    fn gen_bf16(&mut self) -> bf16 {
        let wrapper: Bf16Wrapper = self.gen();
        wrapper.0
    }
    
    fn gen_f16(&mut self) -> f16 {
        let wrapper: F16Wrapper = self.gen();
        wrapper.0
    }
    
    fn gen_range_bf16(&mut self, min: f32, max: f32) -> bf16 {
        rand_uniform_bf16(self, min, max)
    }
    
    fn gen_range_f16(&mut self, min: f32, max: f32) -> f16 {
        rand_uniform_f16(self, min, max)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_bf16_standard_distribution() {
        let mut rng = rand::thread_rng();
        let val: Bf16Wrapper = rng.gen();
        assert!(f32::from(val.0) >= 0.0 && f32::from(val.0) < 1.0);
    }
    
    #[test]
    fn test_f16_standard_distribution() {
        let mut rng = rand::thread_rng();
        let val: F16Wrapper = rng.gen();
        assert!(f32::from(val.0) >= 0.0 && f32::from(val.0) < 1.0);
    }
    
    #[test]
    fn test_rand_uniform_bf16() {
        let mut rng = rand::thread_rng();
        for _ in 0..100 {
            let val = rand_uniform_bf16(&mut rng, -1.0, 1.0);
            assert!(f32::from(val) >= -1.0 && f32::from(val) <= 1.0);
        }
    }
    
    #[test]
    fn test_rand_uniform_f16() {
        let mut rng = rand::thread_rng();
        for _ in 0..100 {
            let val = rand_uniform_f16(&mut rng, -1.0, 1.0);
            assert!(f32::from(val) >= -1.0 && f32::from(val) <= 1.0);
        }
    }
    
    #[test]
    fn test_half_rng_ext() {
        let mut rng = rand::thread_rng();
        let bf16_val = rng.gen_bf16();
        let f16_val = rng.gen_f16();
        
        assert!(f32::from(bf16_val).is_finite());
        assert!(f32::from(f16_val).is_finite());
        
        let range_bf16 = rng.gen_range_bf16(-1.0, 1.0);
        let range_f16 = rng.gen_range_f16(-1.0, 1.0);
        
        assert!(f32::from(range_bf16) >= -1.0 && f32::from(range_bf16) <= 1.0);
        assert!(f32::from(range_f16) >= -1.0 && f32::from(range_f16) <= 1.0);
    }
}