Use BF16 on metal when possible. (#2378)

1 files changed, 16 insertions, 0 deletions

diff --git a/candle-core/src/device.rs b/candle-core/src/device.rs
index 1cd26167..91e56937 100644
--- a/candle-core/src/device.rs
+++ b/candle-core/src/device.rs
@@ -171,6 +171,22 @@ impl Device {
         matches!(self, Self::Metal(_))
     }
 
+    pub fn supports_bf16(&self) -> bool {
+        match self {
+            Self::Cuda(_) | Self::Metal(_) => true,
+            Self::Cpu => false,
+        }
+    }
+
+    /// Return `BF16` for devices that support it, otherwise default to `F32`.
+    pub fn bf16_default_to_f32(&self) -> DType {
+        if self.supports_bf16() {
+            DType::BF16
+        } else {
+            DType::F32
+        }
+    }
+
     pub fn cuda_if_available(ordinal: usize) -> Result<Self> {
         if crate::utils::cuda_is_available() {
             Self::new_cuda(ordinal)