From 45487206bda45355e7e412f964debbbcb5eda483 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 3 Mar 2026 16:10:25 +0100 Subject: [PATCH] fix(ml): allow BF16/F16 in Mamba2 scalar_tensor helper Previously scalar_tensor() rejected BF16 and F16 with an error, blocking BF16 training for Mamba2. Now BF16/F16 scalars are created as F32 then cast to the target dtype, matching Candle's internal precision requirements while supporting half-precision training. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/mamba/mod.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/mamba/mod.rs b/crates/ml/src/mamba/mod.rs index 9f53b1030..c4ddee06e 100644 --- a/crates/ml/src/mamba/mod.rs +++ b/crates/ml/src/mamba/mod.rs @@ -602,7 +602,16 @@ impl Mamba2SSM { reason: e.to_string(), }) }, - DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 | DType::BF16 | DType::F16 => { + DType::BF16 | DType::F16 => { + // Create as F32 scalar, then cast to the target half-precision dtype. + Tensor::new(&[value as f32], device) + .and_then(|t| t.to_dtype(dtype)) + .map_err(|e| MLError::TensorCreationError { + operation: format!("scalar_tensor ({:?} via F32)", dtype), + reason: e.to_string(), + }) + }, + DType::F8E4M3 | DType::U8 | DType::U32 | DType::I64 => { Err(MLError::ModelError(format!( "Unsupported dtype: {:?}", dtype