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

fix(ml): add ensure_training_dtype boundary casts to all model forward methods

Browse Source 
After converting all VarBuilder sites from F32 to training_dtype (BF16 on
Ampere+ GPUs), input tensors from callers remain F32, causing dtype
mismatches in matmul/add/mul operations. This commit adds systematic
boundary casts across all 10 model architectures:

- Input boundary: ensure_training_dtype() at each model forward() entry
- Output boundary: to_dtype(F32) at each model forward() exit
- Internal intermediates: hidden state init, gradient extraction,
  positional encodings, causal masks, SSM state matrices, B-spline
  basis values all cast to match computation dtype
- Ensemble adapters: ensure_training_dtype after Tensor::from_vec

36 files, +293/-72 lines. 2640 tests pass, 0 failures, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-03 17:42:41 +01:00
parent dc8ee1f630
commit 864808e056
36 changed files with 293 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/ml/src/diffusion/denoiser.rs
View File

@@ -63,6 +63,10 @@ impl TimeEmbedding {
let emb = Tensor::cat(&[&sin_emb, &cos_emb], 1)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// Cast to training dtype before projection through BF16 weights
let emb = crate::dqn::mixed_precision::ensure_training_dtype(&emb)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// Project to hidden_dim
self.proj.forward(&emb)
.map_err(|e| MLError::ModelError(e.to_string()))
@@ -201,13 +205,15 @@ impl Denoiser {
///
/// Input x: (batch, data_dim), t: (batch,) → Output: (batch, data_dim)
pub fn forward(&self, x: &Tensor, t: &Tensor) -> Result<Tensor, MLError> {
let x = crate::dqn::mixed_precision::ensure_training_dtype(x)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string());
// Time embedding
let t_emb = self.time_embed.forward(t, &self.device)?;
// Input projection → SiLU
let mut h = self.input_proj.forward(x).map_err(map_err)?;
let mut h = self.input_proj.forward(&x).map_err(map_err)?;
let h_sig = candle_nn::ops::sigmoid(&h).map_err(map_err)?;
h = h.mul(&h_sig).map_err(map_err)?;
@@ -217,7 +223,9 @@ impl Denoiser {
}
// Output projection → predicted noise
self.output_proj.forward(&h).map_err(map_err)
let output = self.output_proj.forward(&h).map_err(map_err)?;
// Cast output back to F32 for API compatibility
output.to_dtype(candle_core::DType::F32).map_err(map_err)
}
}

2
crates/ml/src/diffusion/trainable.rs
View File

@@ -113,6 +113,8 @@ impl UnifiedTrainable for DiffusionTrainableAdapter {
/// Input can be (batch, seq_len * feature_dim) or (batch, seq_len, feature_dim).
/// Returns predicted noise with same shape as flattened input.
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let map_err = |e: candle_core::Error| MLError::ModelError(e.to_string());
// Flatten to (batch, data_dim)

11
crates/ml/src/dqn/distributional_dueling.rs
View File

@@ -39,7 +39,7 @@
//! - Dueling Network Architectures (Wang et al., 2016)
//! - A Distributional Perspective on Reinforcement Learning (Bellemare et al., 2017)
use candle_core::{DType, Device, Tensor};
use candle_core::{Device, Tensor};
use candle_nn::{Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
@@ -228,6 +228,8 @@ impl DistributionalDuelingQNetwork {
/// - A(s,a,z_i): Advantage distribution per action (probability of atom z_i)
/// - mean(A(s,·,z_i)): Mean advantage across actions (ensures identifiability)
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let batch_size = state
.dim(0)
.map_err(|e| MLError::ModelError(format!("Failed to get batch size: {}", e)))?;
@@ -325,6 +327,11 @@ impl DistributionalDuelingQNetwork {
MLError::ModelError(format!("Softmax over atoms failed: {}", e))
})?;
// Cast output back to F32 for API compatibility
let z_probs = z_probs.to_dtype(candle_core::DType::F32).map_err(|e| {
MLError::ModelError(format!("Output dtype cast failed: {}", e))
})?;
Ok(z_probs)
}
@@ -375,7 +382,7 @@ impl DistributionalDuelingQNetwork {
#[cfg(test)]
mod tests {
use super::*;
use candle_core::Device;
use candle_core::{DType, Device};
#[test]
fn test_distributional_dueling_creation() -> anyhow::Result<()> {

22
crates/ml/src/dqn/dqn.rs
View File

@@ -723,6 +723,8 @@ impl Sequential {
/// Forward pass through network
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let mut x = input.clone();
if self.use_noisy_nets {
@@ -755,6 +757,10 @@ impl Sequential {
}
}
// Cast output back to F32 for API compatibility
let x = x.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?;
Ok(x)
}
@@ -1103,7 +1109,9 @@ impl DQN {
/// 2. Dueling only
/// 3. Standard Q-network (fallback)
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
// Auto-convert input to correct device (Candle optimizes if already on correct device)
// Auto-convert input to correct device and dtype
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let state = state
.to_device(&self.device)
.map_err(|e| MLError::ModelError(format!("Failed to move tensor to device: {}", e)))?;
@@ -1129,10 +1137,11 @@ impl DQN {
let num_actions = self.config.num_actions;
// Broadcast atoms to [batch, num_actions, num_atoms]
// Note: atoms stay F32 to match z_probs output (already cast to F32 by sub-networks)
let atoms_tensor = Tensor::from_vec(atoms, num_atoms, &self.device)?
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, num_actions, num_atoms))?;
.unsqueeze(0)?
.unsqueeze(0)?
.broadcast_as((batch_size, num_actions, num_atoms))?;
// Compute expected Q-values: Q(s,a) = Σ(z_i * p_i)
let q_values_dist = (z_probs * atoms_tensor)?.sum(2)?; // Sum over atoms
@@ -1215,6 +1224,9 @@ impl DQN {
q_values
};
// Cast output back to F32 for API compatibility (callers expect f32 tensors)
let q_values = q_values.to_dtype(candle_core::DType::F32)?;
Ok(q_values)
}
@@ -1424,6 +1436,8 @@ impl DQN {
/// the intermediate representation needed by IQN.
fn get_state_embedding(&self, states: &Tensor) -> Result<Tensor, MLError> {
let states = states.to_device(&self.device)?;
let states = crate::dqn::mixed_precision::ensure_training_dtype(&states)
.map_err(|e| MLError::ModelError(e.to_string()))?;
if self.q_network.use_noisy_nets {
let mut x = states;

11
crates/ml/src/dqn/dueling.rs
View File

@@ -30,7 +30,7 @@
//! - A(s,a): Advantage function (per-action)
//! - mean(A(s,·)): Average advantage across all actions (ensures zero mean)
use candle_core::{DType, Device, Tensor};
use candle_core::{Device, Tensor};
use candle_nn::{Linear, Module, VarBuilder, VarMap};
use serde::{Deserialize, Serialize};
@@ -206,6 +206,8 @@ impl DuelingQNetwork {
/// - A(s,a): Advantage per action
/// - mean(A(s,·)): Mean advantage (ensures identifiability)
pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// Shared feature extraction
let mut h = state.clone();
for (i, layer) in self.shared_layers.iter().enumerate() {
@@ -276,6 +278,11 @@ impl DuelingQNetwork {
MLError::ModelError(format!("Q-value combination failed: {}", e))
})?;
// Cast output back to F32 for API compatibility
let q_values = q_values.to_dtype(candle_core::DType::F32).map_err(|e| {
MLError::ModelError(format!("Output dtype cast failed: {}", e))
})?;
Ok(q_values)
}
@@ -323,7 +330,7 @@ impl DuelingQNetwork {
#[cfg(test)]
mod tests {
use super::*;
use candle_core::Device;
use candle_core::{DType, Device};
#[test]
fn test_dueling_network_creation() -> anyhow::Result<()> {

17
crates/ml/src/dqn/mixed_precision.rs
View File

@@ -215,7 +215,22 @@ pub fn training_dtype(device: &candle_core::Device) -> candle_core::DType {
candle_core::DType::F32
}
}
_ => candle_core::DType::F32,
candle_core::Device::Cpu | candle_core::Device::Metal(_) => candle_core::DType::F32,
}
}
/// Cast a tensor to the training dtype for its device if needed.
///
/// This is the canonical "boundary cast" — call it at any model entry
/// point where the input tensor may arrive as F32 from external code
/// (tests, ensemble adapters, raw inference). When dtypes already
/// match the call is essentially free (Candle returns a clone).
pub fn ensure_training_dtype(tensor: &candle_core::Tensor) -> Result<candle_core::Tensor, candle_core::Error> {
let target = training_dtype(tensor.device());
if tensor.dtype() == target {
Ok(tensor.clone())
} else {
tensor.to_dtype(target)
}
}

5
crates/ml/src/dqn/network.rs
View File

@@ -190,6 +190,7 @@ impl NetworkLayers {
xs: &Tensor,
mixed_precision: &Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
) -> CandleResult<Tensor> {
let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?;
let (x_input, use_amp) = match mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
@@ -222,6 +223,7 @@ impl NetworkLayers {
impl Module for NetworkLayers {
fn forward(&self, xs: &Tensor) -> CandleResult<Tensor> {
let xs = crate::dqn::mixed_precision::ensure_training_dtype(xs)?;
let mut x = xs.clone();
// Forward through hidden layers with LeakyReLU activation and dropout
@@ -236,6 +238,9 @@ impl Module for NetworkLayers {
}
}
// Cast output back to F32 for API compatibility
x = x.to_dtype(DType::F32)?;
Ok(x)
}
}

24
crates/ml/src/dqn/quantile_regression.rs
View File

@@ -144,12 +144,16 @@ impl QuantileNetwork {
/// # Returns
/// Quantile values [batch, num_actions, num_quantiles]
pub fn forward(&self, state_embed: &Tensor, taus: &Tensor) -> Result<Tensor, MLError> {
let state_embed = crate::dqn::mixed_precision::ensure_training_dtype(state_embed)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let taus = crate::dqn::mixed_precision::ensure_training_dtype(taus)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let batch_size = state_embed.dim(0)?;
let num_quantiles = taus.dim(1)?;
let num_actions = self.config.num_actions;
// 1. Compute cosine embedding: ψ(τ) = [cos(πi·τ) for i in 1..embedding_dim]
let cos_embed = self.cosine_embedding(taus)?; // [batch, num_quantiles, embedding_dim]
let cos_embed = self.cosine_embedding(&taus)?; // [batch, num_quantiles, embedding_dim]
// 2. State embedding broadcast to match quantile dimension
// [batch, state_dim] → [batch, 1, state_dim] → [batch, num_quantiles, state_dim]
@@ -174,10 +178,14 @@ impl QuantileNetwork {
.map_err(|e| MLError::ModelError(format!("Output layer forward failed: {}", e)))?;
// 7. Transpose to [batch, num_actions, num_quantiles]
quantile_values
let output = quantile_values
.reshape((batch_size, num_quantiles, num_actions))?
.transpose(1, 2)
.map_err(|e| MLError::ModelError(format!("Transpose failed: {}", e)))
.map_err(|e| MLError::ModelError(format!("Transpose failed: {}", e)))?;
// Cast output back to F32 for API compatibility
output.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))
}
/// Cosine embedding for quantile fractions τ
@@ -194,17 +202,18 @@ impl QuantileNetwork {
/// Cosine embeddings [batch, num_quantiles, embedding_dim]
fn cosine_embedding(&self, taus: &Tensor) -> CandleResult<Tensor> {
let device = taus.device();
let target_dtype = taus.dtype();
let batch_size = taus.dim(0)?;
let num_quantiles = taus.dim(1)?;
let embed_dim = self.config.quantile_embedding_dim;
// Create indices [1, 2, 3, ..., embedding_dim]
// Create indices [1, 2, 3, ..., embedding_dim] in same dtype as taus
let indices: Vec<f32> = (1..=embed_dim).map(|i| i as f32).collect();
let indices_tensor = Tensor::from_vec(
indices,
(1, 1, embed_dim),
device,
)?;
)?.to_dtype(target_dtype)?;
// Broadcast taus to [batch, num_quantiles, 1]
let taus_broadcast = taus.unsqueeze(2)?;
@@ -213,8 +222,9 @@ impl QuantileNetwork {
let indices_broadcast = indices_tensor.broadcast_as((batch_size, num_quantiles, embed_dim))?;
let taus_full = taus_broadcast.broadcast_as((batch_size, num_quantiles, embed_dim))?;
// Compute π·i·τ
let pi_tensor = Tensor::full(PI, (batch_size, num_quantiles, embed_dim), device)?;
// Compute π·i·τ (pi_tensor must match taus dtype)
let pi_tensor = Tensor::full(PI, (batch_size, num_quantiles, embed_dim), device)?
.to_dtype(target_dtype)?;
let angles = (pi_tensor * indices_broadcast)? * taus_full;
// cos(π·i·τ)

4
crates/ml/src/ensemble/adapters/diffusion.rs
View File

@@ -118,6 +118,8 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter {
let padded = self.pad_features(&features.values);
let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device)
.map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("Diffusion input dtype cast: {e}")))?;
// Timestep t=1 (minimal noise level for feature processing)
let t = Tensor::from_vec(vec![1_u32], (1,), &self.device)
@@ -160,6 +162,8 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter {
let padded = self.pad_features(&features.values);
let input = Tensor::from_vec(padded, (1, self.data_dim), &self.device)
.map_err(|e| MLError::ModelError(format!("Diffusion input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("Diffusion input dtype cast: {e}")))?;
// Timestep t=1 (minimal noise level for feature processing)
let t = Tensor::from_vec(vec![1_u32], (1,), &self.device)

5
crates/ml/src/ensemble/adapters/dqn.rs
View File

@@ -8,6 +8,7 @@ use std::sync::Mutex;
use candle_core::{Device, Tensor};
use crate::dqn::dqn::{DQNConfig, DQN};
use crate::dqn::mixed_precision::ensure_training_dtype;
use crate::ensemble::inference_adapter::{
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
};
@@ -70,9 +71,11 @@ impl ModelInferenceAdapter for DqnInferenceAdapter {
let f32_values: Vec<f32> = features.values.iter().map(|&v| v as f32).collect();
let len = f32_values.len();
// Create input tensor [1, feature_dim]
// Create input tensor [1, feature_dim] and cast to training dtype
let input = Tensor::from_vec(f32_values, (1, len), &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to create input tensor: {e}")))?;
let input = ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("Input dtype cast failed: {e}")))?;
// Run forward pass through the Q-network
let model = self

4
crates/ml/src/ensemble/adapters/kan.rs
View File

@@ -99,6 +99,8 @@ impl ModelInferenceAdapter for KanInferenceAdapter {
let padded = self.pad_features(&features.values);
let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device)
.map_err(|e| MLError::ModelError(format!("Failed to create KAN input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("KAN input dtype cast: {e}")))?;
let model = self
.model
@@ -139,6 +141,8 @@ impl ModelInferenceAdapter for KanInferenceAdapter {
let padded = self.pad_features(&features.values);
let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device)
.map_err(|e| MLError::ModelError(format!("KAN input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("KAN input dtype cast: {e}")))?;
let model = self
.model

10
crates/ml/src/ensemble/adapters/mamba2.rs
View File

@@ -9,8 +9,9 @@
use std::collections::VecDeque;
use std::sync::Mutex;
use candle_core::{DType, Device, Tensor};
use candle_core::{Device, Tensor};
use crate::dqn::mixed_precision::ensure_training_dtype;
use crate::ensemble::inference_adapter::{
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
};
@@ -148,10 +149,9 @@ impl ModelInferenceAdapter for Mamba2InferenceAdapter {
)
.map_err(|e| MLError::ModelError(format!("Failed to create Mamba2 input tensor: {e}")))?;
// Cast to F32 (Mamba2 uses DType::F32 for GPU throughput)
let input = input
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("Failed to cast input to F32: {e}")))?;
// Cast to training dtype (BF16 on Ampere+, F32 on CPU/older GPUs)
let input = ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("Failed to cast input to training dtype: {e}")))?;
// Run forward pass (needs &mut self)
let mut model = self

8
crates/ml/src/ensemble/adapters/tft.rs
View File

@@ -247,11 +247,13 @@ impl ModelInferenceAdapter for TftInferenceAdapter {
// Future (known) features: calendar-derived
let future_f32 = self.build_future_features(last_ts);
// Build tensors
// Build tensors and cast to training dtype
let static_tensor =
Tensor::from_vec(static_f32, (1, self.num_static), &self.device).map_err(|e| {
MLError::ModelError(format!("Failed to create static tensor: {e}"))
})?;
let static_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&static_tensor)
.map_err(|e| MLError::ModelError(format!("Static dtype cast: {e}")))?;
let hist_tensor = Tensor::from_vec(
hist_f32,
@@ -259,6 +261,8 @@ impl ModelInferenceAdapter for TftInferenceAdapter {
&self.device,
)
.map_err(|e| MLError::ModelError(format!("Failed to create historical tensor: {e}")))?;
let hist_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&hist_tensor)
.map_err(|e| MLError::ModelError(format!("Historical dtype cast: {e}")))?;
let future_tensor = Tensor::from_vec(
future_f32,
@@ -266,6 +270,8 @@ impl ModelInferenceAdapter for TftInferenceAdapter {
&self.device,
)
.map_err(|e| MLError::ModelError(format!("Failed to create future tensor: {e}")))?;
let future_tensor = crate::dqn::mixed_precision::ensure_training_dtype(&future_tensor)
.map_err(|e| MLError::ModelError(format!("Future dtype cast: {e}")))?;
// --- forward pass (model lock) ---
let quantile_preds = {

10
crates/ml/src/ensemble/adapters/tggn.rs
View File

@@ -31,16 +31,20 @@ impl TggnProjection {
}
fn forward(&self, input: &Tensor) -> MLResult<Tensor> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(format!("TGGN input dtype cast: {e}")))?;
let h = self
.linear1
.forward(input)
.forward(&input)
.map_err(|e| MLError::ModelError(format!("TGGN forward linear1: {e}")))?;
let h = h
.relu()
.map_err(|e| MLError::ModelError(format!("TGGN forward relu: {e}")))?;
self.linear2
let out = self.linear2
.forward(&h)
.map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}")))
.map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}")))?;
out.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("TGGN output dtype cast: {e}")))
}
}

10
crates/ml/src/ensemble/adapters/tlob.rs
View File

@@ -43,9 +43,11 @@ impl TlobProjection {
}
fn forward(&self, input: &Tensor) -> MLResult<Tensor> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(format!("TLOB input dtype cast: {e}")))?;
let h = self
.linear1
.forward(input)
.forward(&input)
.map_err(|e| MLError::ModelError(format!("TLOB forward linear1: {e}")))?;
let h = h
.relu()
@@ -57,9 +59,11 @@ impl TlobProjection {
let h = h
.relu()
.map_err(|e| MLError::ModelError(format!("TLOB forward relu2: {e}")))?;
self.linear3
let out = self.linear3
.forward(&h)
.map_err(|e| MLError::ModelError(format!("TLOB forward linear3: {e}")))
.map_err(|e| MLError::ModelError(format!("TLOB forward linear3: {e}")))?;
out.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("TLOB output dtype cast: {e}")))
}
}

4
crates/ml/src/ensemble/adapters/xlstm.rs
View File

@@ -166,6 +166,8 @@ impl ModelInferenceAdapter for XlstmInferenceAdapter {
&self.device,
)
.map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("xLSTM input dtype cast: {e}")))?;
// XLSTMNetwork::forward takes &self (not &mut self)
let model = self
@@ -247,6 +249,8 @@ impl ModelInferenceAdapter for XlstmInferenceAdapter {
&self.device,
)
.map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?;
let input = crate::dqn::mixed_precision::ensure_training_dtype(&input)
.map_err(|e| MLError::ModelError(format!("xLSTM input dtype cast: {e}")))?;
let model = self
.model

2
crates/ml/src/gradient_utils.rs
View File

@@ -44,7 +44,7 @@ pub fn clip_grad_norm(
// First pass: Calculate the total L2 norm of all gradients
for var in vars.iter() {
if let Some(grad) = grads.get(var) {
let norm_sq = grad.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
let norm_sq = grad.sqr()?.sum_all()?.to_dtype(candle_core::DType::F32)?.to_scalar::<f32>()? as f64;
total_norm_sq += norm_sq;
}
}

14
crates/ml/src/kan/layer.rs
View File

@@ -109,13 +109,25 @@ impl KANLayer {
MLError::ModelError(format!("KAN basis reshape failed: {}", e))
})?;
// Cast basis_flat to match coefficients dtype (BSpline evaluates in F32
// for precision, but coefficients are in training dtype from VarBuilder)
let coeff_dtype = self.coefficients.dtype();
let basis_flat = basis_flat.to_dtype(coeff_dtype).map_err(|e| {
MLError::ModelError(format!("KAN basis dtype cast failed: {}", e))
})?;
// Spline contribution: (batch, in_dim * num_bases) @ (in_dim * num_bases, out_dim)
let spline_out = basis_flat.matmul(&self.coefficients).map_err(|e| {
MLError::ModelError(format!("KAN spline matmul failed: {}", e))
})?;
// Cast input to match residual_weight dtype for matmul compatibility
let input_cast = input.to_dtype(coeff_dtype).map_err(|e| {
MLError::ModelError(format!("KAN input dtype cast failed: {}", e))
})?;
// Residual contribution: (batch, in_dim) @ (in_dim, out_dim)
let residual_out = input.matmul(&self.residual_weight).map_err(|e| {
let residual_out = input_cast.matmul(&self.residual_weight).map_err(|e| {
MLError::ModelError(format!("KAN residual matmul failed: {}", e))
})?;

6
crates/ml/src/kan/network.rs
View File

@@ -62,11 +62,15 @@ impl KANNetwork {
/// Input shape: `(batch, layer_widths[0])`
/// Output shape: `(batch, layer_widths[last])`
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let mut x = input.clone();
for layer in &self.layers {
x = layer.forward(&x)?;
}
Ok(x)
// Cast output back to F32 for API compatibility
x.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))
}
}

8
crates/ml/src/kan/spline.rs
View File

@@ -164,6 +164,10 @@ impl BSplineBasis {
let x_flat = x.flatten_all().map_err(|e| {
MLError::ModelError(format!("BSpline flatten failed: {e}"))
})?;
// BSpline evaluation requires F32 precision for index computation
let x_flat = x_flat.to_dtype(DType::F32).map_err(|e| {
MLError::ModelError(format!("BSpline dtype cast: {e}"))
})?;
let n = x_flat.dims1().map_err(|e| {
MLError::ModelError(format!("BSpline dims failed: {e}"))
})?;
@@ -200,6 +204,10 @@ impl BSplineBasis {
let x_flat = x.flatten_all().map_err(|e| {
MLError::ModelError(format!("BSpline flatten failed: {e}"))
})?;
// BSpline evaluation uses floor/ceil/frac which require F32 precision
let x_flat = x_flat.to_dtype(DType::F32).map_err(|e| {
MLError::ModelError(format!("BSpline dtype cast: {e}"))
})?;
let n = x_flat.dims1().map_err(|e| {
MLError::ModelError(format!("BSpline dims failed: {e}"))
})?;

4
crates/ml/src/kan/trainable.rs
View File

@@ -95,7 +95,9 @@ impl UnifiedTrainable for KANTrainableAdapter {
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
self.network.forward(input)
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
self.network.forward(&input)
}
fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError> {

13
crates/ml/src/liquid/candle_cfc.rs
View File

@@ -4,7 +4,7 @@
//! This is the training path; the existing FixedPoint implementation in cells.rs/network.rs
//! remains the production inference path.
use candle_core::{DType, Tensor};
use candle_core::Tensor;
use candle_nn::{Linear, Module, VarBuilder};
use serde::{Deserialize, Serialize};
@@ -121,6 +121,8 @@ impl BackboneMLP {
/// - `f_out` is tanh-activated (bounded in [-1, 1])
/// - `tau_out` is raw (sigmoid scaling applied in the CfC cell)
pub fn forward(&self, input: &Tensor) -> Result<(Tensor, Tensor), MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::InferenceError(e.to_string()))?;
let mut x = input.clone();
for layer in &self.layers {
x = layer
@@ -285,7 +287,7 @@ impl CandleCfCNetwork {
let mut h = Tensor::zeros(
(batch_size, self.config.hidden_size),
DType::F32,
crate::dqn::mixed_precision::training_dtype(device),
device,
)
.map_err(|e| MLError::InferenceError(format!("CfC init hidden: {}", e)))?;
@@ -308,7 +310,12 @@ impl CandleCfCNetwork {
/// Forward compatible with UnifiedTrainable (3D input, default dt=0.01)
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
self.forward_sequence(input, 0.01)
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::InferenceError(e.to_string()))?;
let output = self.forward_sequence(&input, 0.01)?;
// Cast output back to F32 for API compatibility
output.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::InferenceError(format!("Output dtype cast failed: {}", e)))
}
/// Approximate parameter count for reporting purposes

3
crates/ml/src/liquid/training.rs
View File

@@ -462,7 +462,7 @@ impl LiquidTrainer {
// === Candle-based CfC v2 Trainer (2026 modernization) ===
use candle_core::{DType, Device, Tensor};
use candle_core::{Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use super::candle_cfc::{CandleCfCNetwork, CfCTrainConfig};
@@ -663,6 +663,7 @@ impl TrainingUtils {
#[cfg(test)]
mod tests {
use super::*;
use candle_core::DType;
use crate::liquid::candle_cfc::DeviceConfig;
#[test]

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

@@ -779,6 +779,8 @@ impl Mamba2SSM {
/// - Tensor operations fail
#[instrument(skip(self, input))]
pub fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let start = Instant::now();
// OPTIMIZATION: Device affinity check (catch cross-device transfers early)
@@ -792,7 +794,7 @@ impl Mamba2SSM {
}
// Input projection
let mut hidden = self.input_projection.forward(input)?;
let mut hidden = self.input_projection.forward(&input)?;
// Process through each layer - collect indices first to avoid borrow conflicts
let num_layers = self.ssd_layers.len();
@@ -818,6 +820,10 @@ impl Mamba2SSM {
let output_raw = self.output_projection.forward(&hidden)?;
let output = crate::cuda_compat::manual_sigmoid(&output_raw)?;
// Cast output back to F32 for API compatibility
let output = output.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?;
// OPTIMIZATION: Update performance metrics with VecDeque (O(1) instead of O(n))
let inference_time = start.elapsed();
self.total_inferences.fetch_add(1, Ordering::Relaxed);
@@ -885,7 +891,9 @@ impl Mamba2SSM {
C.dims()
);
let batch_size = scanned_states.dim(0)?;
let C_t = C.t()?.contiguous()?;
// Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16)
let C_cast = C.to_dtype(scanned_states.dtype())?;
let C_t = C_cast.t()?.contiguous()?;
let C_broadcasted =
C_t.unsqueeze(0)?
.broadcast_as((batch_size, C_t.dim(0)?, C_t.dim(1)?))?;
@@ -963,7 +971,9 @@ impl Mamba2SSM {
self.config.d_state
);
let B_t = B.t()?.contiguous()?;
// Cast B to match input dtype (SSM state is F32 but input may be BF16)
let B_cast = B.to_dtype(input.dtype())?;
let B_t = B_cast.t()?.contiguous()?;
let d_inner = B_t.dim(0)?;
let d_state = B_t.dim(1)?;
let B_broadcasted = B_t
@@ -1684,7 +1694,9 @@ impl Mamba2SSM {
// scanned_states: [32, 60, 16]
// C stored as: [d_inner, d_state] = [512, 16]
// Need: [batch, d_state, d_inner] = [32, 16, 512]
let C_t = C.t()?.contiguous()?; // [512, 16] → [16, 512]
// Cast C to match scanned_states dtype (SSM state is F32 but computation may be BF16)
let C_cast = C.to_dtype(scanned_states.dtype())?;
let C_t = C_cast.t()?.contiguous()?; // [512, 16] → [16, 512]
trace!("C transposed (d_state, d_inner): {:?}", C_t.dims());
// Now broadcast [16, 512] to [32, 16, 512]
@@ -1744,6 +1756,9 @@ impl Mamba2SSM {
let mut result = Tensor::zeros((batch_size, seq_len, d_state), input.dtype(), device)?;
let mut current_state = Tensor::zeros((batch_size, d_state), input.dtype(), device)?;
// Cast A to match input dtype (SSM state matrices are F32 but computation may be BF16)
let A_cast = A.to_dtype(input.dtype())?;
// Sequential scan with state transitions (maintaining gradients)
for t in 0..seq_len {
let x_t = input.narrow(1, t, 1)?.squeeze(1)?;
@@ -1752,7 +1767,7 @@ impl Mamba2SSM {
// State transition: h_t = h_{t-1} @ A^T + x_t
// current_state [batch, d_state] × A.t() [d_state, d_state] = [batch, d_state]
// This is the correct way to do batch SSM state transitions
current_state = (current_state.matmul(&A.t()?)? + &x_t)?;
current_state = (current_state.matmul(&A_cast.t()?)? + &x_t)?;
// Write directly to result tensor (no Vec accumulation, no Tensor::cat doubling)
let current_unsqueezed = current_state.unsqueeze(1)?;
@@ -1838,7 +1853,9 @@ impl Mamba2SSM {
// input: [batch, seq, d_inner], B: [d_state, d_inner]
// B.t(): [d_inner, d_state] → explicit repeat to [batch, d_inner, d_state]
let batch_size = input.dim(0)?;
let B_t = B.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state]
// Cast B to match input dtype (SSM state is F32 but input may be BF16)
let B_cast = B.to_dtype(input.dtype())?;
let B_t = B_cast.t()?.contiguous()?; // [d_state, d_inner] → [d_inner, d_state]
// CRITICAL FIX: Use repeat/expand instead of broadcast_as for CUDA compatibility
// Create [batch, d_inner, d_state] by repeating the [d_inner, d_state] tensor

8
crates/ml/src/mamba/scan_algorithms.rs
View File

@@ -327,9 +327,11 @@ impl ParallelScanEngine {
let alpha_fp = FixedPoint::from_f64(0.9); // Decay factor
let beta_fp = FixedPoint::from_f64(0.1); // Input weight
// Use F32 to match state/input tensor dtype (model-wide F32 for GPU throughput)
let alpha = Tensor::full(alpha_fp.to_f64() as f32, state.shape(), state.device())?;
let beta = Tensor::full(beta_fp.to_f64() as f32, input.shape(), input.device())?;
// Match state/input tensor dtype (may be BF16 during mixed-precision training)
let alpha = Tensor::full(alpha_fp.to_f64() as f32, state.shape(), state.device())?
.to_dtype(state.dtype())?;
let beta = Tensor::full(beta_fp.to_f64() as f32, input.shape(), input.device())?
.to_dtype(input.dtype())?;
let decayed_state = (state * alpha)?;
let input_contribution = (input * beta)?;

24
crates/ml/src/ppo/lstm_networks.rs
View File

@@ -106,10 +106,12 @@ impl LSTMPolicyNetwork {
h_t: &Tensor,
c_t: &Tensor,
) -> Result<(Tensor, Tensor, Tensor), MLError> {
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// Project input: [batch, input_dim] → [batch, hidden_dim]
let x = self
.input_layer
.forward(state)
.forward(&state)
.map_err(|e| MLError::ModelError(format!("Input layer forward failed: {}", e)))?;
// Apply ReLU activation
@@ -160,6 +162,14 @@ impl LSTMPolicyNetwork {
.forward(&x)
.map_err(|e| MLError::ModelError(format!("Output layer forward failed: {}", e)))?;
// Cast outputs back to F32 for API compatibility
let logits = logits.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Logits dtype cast failed: {}", e)))?;
let new_h = new_h.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Hidden state dtype cast failed: {}", e)))?;
let new_c = new_c.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Cell state dtype cast failed: {}", e)))?;
Ok((logits, new_h, new_c))
}
@@ -320,10 +330,12 @@ impl LSTMValueNetwork {
h_t: &Tensor,
c_t: &Tensor,
) -> Result<(Tensor, Tensor, Tensor), MLError> {
let state = crate::dqn::mixed_precision::ensure_training_dtype(state)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// Project input: [batch, input_dim] → [batch, hidden_dim]
let x = self
.input_layer
.forward(state)
.forward(&state)
.map_err(|e| MLError::ModelError(format!("Input layer forward failed: {}", e)))?;
// Apply ReLU activation
@@ -379,6 +391,14 @@ impl LSTMValueNetwork {
.squeeze(1)
.map_err(|e| MLError::ModelError(format!("Failed to squeeze value output: {}", e)))?;
// Cast outputs back to F32 for API compatibility
let value = value.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Value dtype cast failed: {}", e)))?;
let new_h = new_h.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Hidden state dtype cast failed: {}", e)))?;
let new_c = new_c.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Cell state dtype cast failed: {}", e)))?;
Ok((value, new_h, new_c))
}

18
crates/ml/src/ppo/ppo.rs
View File

@@ -427,7 +427,9 @@ impl PolicyNetwork {
input: &Tensor,
mixed_precision: &Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
) -> Result<Tensor, MLError> {
let (mut x, use_amp) = match mixed_precision {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let (mut x, _use_amp) = match mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
match input.to_dtype(target_dtype) {
@@ -450,9 +452,10 @@ impl PolicyNetwork {
}
}
if use_amp {
// Always cast output back to F32 for API compatibility
if x.dtype() != DType::F32 {
x = x.to_dtype(DType::F32).map_err(|e| {
MLError::ModelError(format!("AMP cast back to F32 failed: {}", e))
MLError::ModelError(format!("Output dtype cast to F32 failed: {}", e))
})?;
}
@@ -671,7 +674,9 @@ impl ValueNetwork {
input: &Tensor,
mixed_precision: &Option<crate::dqn::mixed_precision::MixedPrecisionConfig>,
) -> Result<Tensor, MLError> {
let (mut x, use_amp) = match mixed_precision {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let (mut x, _use_amp) = match mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
match input.to_dtype(target_dtype) {
@@ -697,9 +702,10 @@ impl ValueNetwork {
// Squeeze the last dimension (from [batch, 1] to [batch])
x = x.squeeze(1)?;
if use_amp {
// Always cast output back to F32 for API compatibility
if x.dtype() != DType::F32 {
x = x.to_dtype(DType::F32).map_err(|e| {
MLError::ModelError(format!("AMP cast back to F32 failed: {}", e))
MLError::ModelError(format!("Output dtype cast to F32 failed: {}", e))
})?;
}

18
crates/ml/src/tft/mod.rs
View File

@@ -566,10 +566,16 @@ impl TemporalFusionTransformer {
future_features: &Tensor,
use_checkpointing: bool,
) -> Result<Tensor, MLError> {
let static_features = crate::dqn::mixed_precision::ensure_training_dtype(static_features)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let historical_features = crate::dqn::mixed_precision::ensure_training_dtype(historical_features)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let future_features = crate::dqn::mixed_precision::ensure_training_dtype(future_features)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let start_time = Instant::now();
// Validate input dimensions
self.validate_input_dimensions(static_features, historical_features, future_features)?;
self.validate_input_dimensions(&static_features, &historical_features, &future_features)?;
// Log device placement for debugging
debug!("Forward pass device check:");
@@ -580,7 +586,7 @@ impl TemporalFusionTransformer {
// 1. Variable Selection Networks (skip absent feature paths)
let static_encoded = if let Some(ref mut static_vsn) = self.static_variable_selection {
let static_selected = static_vsn.forward(static_features, None)?;
let static_selected = static_vsn.forward(&static_features, None)?;
let encoder = self.static_encoder.as_mut().ok_or_else(|| {
MLError::ModelError(
"static_encoder must exist when static_variable_selection exists".to_owned(),
@@ -597,7 +603,7 @@ impl TemporalFusionTransformer {
let historical_selected = self
.historical_variable_selection
.forward(historical_features, None)?;
.forward(&historical_features, None)?;
let historical_encoded = if use_checkpointing {
self.historical_encoder
@@ -608,7 +614,7 @@ impl TemporalFusionTransformer {
};
let future_encoded = if let Some(ref mut future_vsn) = self.future_variable_selection {
let future_selected = future_vsn.forward(future_features, None)?;
let future_selected = future_vsn.forward(&future_features, None)?;
let encoder = self.future_encoder.as_mut().ok_or_else(|| {
MLError::ModelError(
"future_encoder must exist when future_variable_selection exists".to_owned(),
@@ -661,6 +667,10 @@ impl TemporalFusionTransformer {
debug!(" quantile_preds: {:?}", quantile_preds.device());
// Cast output back to F32 for API compatibility
let quantile_preds = quantile_preds.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("Output dtype cast failed: {}", e)))?;
// Update performance metrics
let latency = start_time.elapsed().as_micros() as u64;
self.update_performance_metrics(latency);

5
crates/ml/src/tft/temporal_attention.rs
View File

@@ -350,14 +350,18 @@ impl TemporalSelfAttention {
// Add positional encoding (lightweight, no checkpointing needed)
let pos_encoding = self.positional_encoding.forward(seq_len)?;
// Cast positional encoding to match input dtype (encoding is F32, input may be BF16)
let pos_encoding = pos_encoding.to_dtype(x.dtype())?;
let pos_encoding_batch = pos_encoding
.unsqueeze(0)?
.broadcast_as((batch_size, seq_len, hidden_dim))?;
let x_with_pos = (x + &pos_encoding_batch)?;
// Create causal mask if needed (now [1, seq_len, seq_len])
// Cast mask to match input dtype (mask is F32 but attention may be BF16)
let mask = if causal_mask {
let base_mask = self.create_causal_mask(seq_len)?;
let base_mask = base_mask.to_dtype(x.dtype())?;
// Broadcast to [batch_size, seq_len, seq_len]
Some(base_mask.broadcast_as((batch_size, seq_len, seq_len))?)
} else {
@@ -393,6 +397,7 @@ impl TemporalSelfAttention {
for (i, weights) in attention_weights.iter().enumerate() {
let mean_weight = weights
.mean_all()
.and_then(|t| t.to_dtype(candle_core::DType::F32))
.and_then(|t| t.to_vec0::<f32>())
.unwrap_or(0.0) as f64;
weight_stats.insert(format!("head_{}_mean", i), mean_weight);

2
crates/ml/src/tft/variable_selection.rs
View File

@@ -130,7 +130,7 @@ impl VariableSelectionNetwork {
fn update_importance_scores(&mut self, attention_weights: &Tensor) -> Result<(), MLError> {
// Compute mean attention weights across batch and time
let mean_weights = attention_weights.mean_keepdim(0)?.mean_keepdim(1)?; // [1, 1, input_size]
let weights_vec = mean_weights.flatten_all()?.to_vec1::<f32>()?;
let weights_vec = mean_weights.flatten_all()?.to_dtype(candle_core::DType::F32)?.to_vec1::<f32>()?;
// Clear previous scores to prevent memory growth (HashMap maintains capacity but releases entries)
self.importance_scores.clear();

9
crates/ml/src/tgnn/trainable_adapter.rs
View File

@@ -139,9 +139,11 @@ impl UnifiedTrainable for TGGNTrainableAdapter {
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// input: [batch, node_dim]
// input_linear: node_dim -> hidden_dim
let hidden = self.input_linear.forward(input).map_err(|e| {
let hidden = self.input_linear.forward(&input).map_err(|e| {
MLError::ModelError(format!("Input linear forward failed: {}", e))
})?;
@@ -155,6 +157,11 @@ impl UnifiedTrainable for TGGNTrainableAdapter {
MLError::ModelError(format!("Output linear forward failed: {}", e))
})?;
// Cast output back to F32 for API compatibility
let output = output.to_dtype(candle_core::DType::F32).map_err(|e| {
MLError::ModelError(format!("Output dtype cast failed: {}", e))
})?;
Ok(output)
}

7
crates/ml/src/tlob/trainable_adapter.rs
View File

@@ -170,6 +170,8 @@ impl UnifiedTrainable for TLOBTrainableAdapter {
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
// input shape: [batch, seq_len, feature_dim] or [batch, seq_len*feature_dim]
// Flatten to [batch, seq_len*feature_dim] if 3D
let dims = input.shape().dims();
@@ -200,6 +202,11 @@ impl UnifiedTrainable for TLOBTrainableAdapter {
MLError::ModelError(format!("Output linear forward failed: {}", e))
})?;
// Cast output back to F32 for API compatibility
let output = output.to_dtype(candle_core::DType::F32).map_err(|e| {
MLError::ModelError(format!("Output dtype cast failed: {}", e))
})?;
Ok(output)
}

7
crates/ml/src/xlstm/mlstm.rs
View File

@@ -12,7 +12,7 @@
//! batch item. With default hidden=64, heads=4: head_dim=16, so C is 4×16×16=1024
//! floats per sample -- very manageable. Be cautious with hidden>128 or heads<2.
use candle_core::{DType, Tensor};
use candle_core::Tensor;
use candle_nn::{linear, Linear, Module, VarBuilder};
use crate::MLError;
@@ -93,9 +93,10 @@ impl MLSTMCell {
let (h_prev, c_flat_prev) = match state {
Some((h, c)) => (h.clone(), c.clone()),
None => {
let h_zeros = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, dev)
let dtype = crate::dqn::mixed_precision::training_dtype(dev);
let h_zeros = Tensor::zeros(&[batch, self.hidden_dim], dtype, dev)
.map_err(|e| MLError::ModelError(format!("mLSTM h_zeros: {e}")))?;
let c_zeros = Tensor::zeros(&[batch, c_flat_size], DType::F32, dev)
let c_zeros = Tensor::zeros(&[batch, c_flat_size], dtype, dev)
.map_err(|e| MLError::ModelError(format!("mLSTM c_zeros: {e}")))?;
(h_zeros, c_zeros)
}

17
crates/ml/src/xlstm/network.rs
View File

@@ -6,7 +6,7 @@
//! Block allocation: blocks 0..N*slstm_ratio are sLSTM, rest are mLSTM.
//! The first block has input_dim → hidden_dim, subsequent blocks are hidden_dim → hidden_dim.
use candle_core::{DType, Tensor};
use candle_core::Tensor;
use candle_nn::{linear, Linear, Module, VarBuilder};
use crate::MLError;
@@ -85,6 +85,8 @@ impl XLSTMNetwork {
///
/// Returns `(batch, output_dim)`.
pub fn forward(&self, input: &Tensor) -> Result<Tensor, MLError> {
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
let dims = input.dims();
let (batch, seq_len, _features) = match dims.len() {
2 => (dims[0], 1_usize, dims[1]),
@@ -101,7 +103,11 @@ impl XLSTMNetwork {
// State per block: (h, c_flat)
let mut states: Vec<Option<(Tensor, Tensor)>> = vec![None; self.blocks.len()];
let mut final_h = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, input.device())
let mut final_h = Tensor::zeros(
&[batch, self.hidden_dim],
crate::dqn::mixed_precision::training_dtype(input.device()),
input.device(),
)
.map_err(|e| MLError::ModelError(format!("xLSTM zeros: {e}")))?;
for t in 0..seq_len {
@@ -139,8 +145,11 @@ impl XLSTMNetwork {
}
// Output projection
self.output_head.forward(&final_h)
.map_err(|e| MLError::ModelError(format!("xLSTM output: {e}")))
let output = self.output_head.forward(&final_h)
.map_err(|e| MLError::ModelError(format!("xLSTM output: {e}")))?;
// Cast output back to F32 for API compatibility
output.to_dtype(candle_core::DType::F32)
.map_err(|e| MLError::ModelError(format!("xLSTM output dtype cast: {e}")))
}
}

8
crates/ml/src/xlstm/slstm.rs
View File

@@ -5,7 +5,7 @@
//! This prevents gradient vanishing over long sequences. A normalizer state `n_t`
//! stabilizes the exponential values.
use candle_core::{DType, Tensor};
use candle_core::Tensor;
use candle_nn::{linear, Linear, Module, VarBuilder};
use crate::MLError;
@@ -57,7 +57,11 @@ impl SLSTMCell {
let (h_prev, c_prev) = match state {
Some((h, c)) => (h.clone(), c.clone()),
None => {
let zeros = Tensor::zeros(&[batch, self.hidden_dim], DType::F32, dev)
let zeros = Tensor::zeros(
&[batch, self.hidden_dim],
crate::dqn::mixed_precision::training_dtype(dev),
dev,
)
.map_err(|e| MLError::ModelError(format!("sLSTM init zeros: {e}")))?;
(zeros.clone(), zeros)
}

4
crates/ml/src/xlstm/trainable.rs
View File

@@ -88,7 +88,9 @@ impl UnifiedTrainable for XLSTMTrainableAdapter {
}
fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
self.network.forward(input)
let input = crate::dqn::mixed_precision::ensure_training_dtype(input)
.map_err(|e| MLError::ModelError(e.to_string()))?;
self.network.forward(&input)
}
fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError> {