huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a hack for generating random uniform/normal for f16/bf16. (#1228)

Browse Source
This commit is contained in:
Laurent Mazare
2023-10-31 21:27:59 +01:00
committed by GitHub
parent c12ad45562
commit 36fb84f038
1 changed files with 16 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
candle-core/src/device.rs
View File

@ -185,11 +185,17 @@ impl Device {
Ok(Storage::Cpu(storage))
}
Device::Cuda(device) => {
// TODO: Remove the special case if we start supporting generating f16/bf16 directly.
if dtype == DType::F16 || dtype == DType::BF16 {
let storage = device.rand_uniform(shape, DType::F32, lo, up)?;
Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
} else {
let storage = device.rand_uniform(shape, dtype, lo, up)?;
Ok(Storage::Cuda(storage))
}
}
}
}
pub(crate) fn rand_uniform<T: crate::FloatDType>(
&self,
@ -213,11 +219,17 @@ impl Device {
Ok(Storage::Cpu(storage))
}
Device::Cuda(device) => {
// TODO: Remove the special case if we start supporting generating f16/bf16 directly.
if dtype == DType::F16 || dtype == DType::BF16 {
let storage = device.rand_normal(shape, DType::F32, mean, std)?;
Storage::Cuda(storage).to_dtype(&crate::Layout::contiguous(shape), dtype)
} else {
let storage = device.rand_normal(shape, dtype, mean, std)?;
Ok(Storage::Cuda(storage))
}
}
}
}
pub(crate) fn rand_normal<T: crate::FloatDType>(
&self,