gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ml): allow BF16/F16 in Mamba2 scalar_tensor helper

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-03 16:10:25 +01:00
parent 868daff02a
commit 45487206bd
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
crates/ml/src/mamba/mod.rs
View File

@@ -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