starkat99 / half-rs

Half-precision floating point types f16 and bf16 for Rust.
https://docs.rs/half/
Other
229 stars 51 forks source link

Implement rand_distr::SampleUniform for f16 #82

Closed ViliamVadocz closed 1 year ago

ViliamVadocz commented 1 year ago

We would like to use f16 in dfdx, a deep learning library. It would be very convenient if this crate implemented rand_distr::SampleUniform for f16.

Relevant issue in dfdx: https://github.com/coreylowman/dfdx/issues/423

The trait would allow us to sample a random f16 from a uniform distribution. This seems like a generally useful feature, perhaps behind a feature gate. Other alternatives exist from our side, but they make the implementation more annoying.

I would be willing to implement this feature if a maintainer agrees that it is a good idea. Closing this issue is also a completely valid response if you believe that this trait implementation does not belong in this crate.

coreylowman commented 1 year ago

We'd really need to sample from any Distribution - the implementation can just sample a f32 and convert that to a f16, so it wouldn't need to do anything complicated.

I wonder if some sort of blanket impl like so might work, but I think it'd run into the foreign types problem:

impl<X: Distribution<f32>> Distribution<f16> for X {
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> f16 {
        f16::from_f32(<Self as Distribution<f32>>::sample(self, rng))
    }
}
coreylowman commented 1 year ago

Okay this is as far as I got:

use rand_distr::Distribution;

macro_rules! impl_distribution_via_f32 {
    ($Distr:ty) => {
        impl Distribution<crate::f16> for $Distr {
            fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> crate::f16 {
                crate::f16::from_f32(<Self as Distribution<f32>>::sample(self, rng))
            }
        }
    };
}

impl_distribution_via_f32!(rand_distr::Standard);
impl_distribution_via_f32!(rand_distr::StandardNormal);
impl_distribution_via_f32!(rand_distr::Exp1);
impl_distribution_via_f32!(rand_distr::Open01);
impl_distribution_via_f32!(rand_distr::OpenClosed01);

impl rand_distr::uniform::SampleUniform for crate::f16 {
    type Sampler = rand_distr::uniform::UniformFloat<Self>;
}

impl rand_distr::uniform::UniformSampler for rand_distr::uniform::UniformFloat<crate::f16> {
    type X = crate::f16;
    fn new<B1, B2>(low: B1, high: B2) -> Self
    where
        B1: rand_distr::uniform::SampleBorrow<Self::X> + Sized,
        B2: rand_distr::uniform::SampleBorrow<Self::X> + Sized,
    {
        todo!()
    }
    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
    where
        B1: rand_distr::uniform::SampleBorrow<Self::X> + Sized,
        B2: rand_distr::uniform::SampleBorrow<Self::X> + Sized,
    {
        todo!()
    }
    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Self::X {
        todo!()
    }
}

However the uniform implementation fails due to foreign traits:

error[E0117]: only traits defined in the current crate can be implemented for types defined outside of the crate
  --> src\sample.rs:23:1
   |
23 | impl rand_distr::uniform::UniformSampler for rand_distr::uniform::UniformFloat<crate::f16> {
   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^---------------------------------------------
   | |                                            |
   | |                                            `UniformFloat` is not defined in the current crate
   | impl doesn't use only types from inside the current crate
   |
   = note: define and implement a trait or new type instead

For more information about this error, try `rustc --explain E0117`.
error: could not compile `half` due to previous error

Maybe I'm missing something about doing custom Uniform distributions?