rustorch 0.6.29

//! f32統一ハイブリッドニューラルネットワークモジュール
//! f32 Unified Hybrid Neural Network Module
//!
//! フェーズ5: 高度ニューラルネットワーク機能
//! Phase 5: Advanced Neural Network Features
//!
//! このモジュールは、f32精度で最適化されたニューラルネットワーク機能を提供します。
//! Neural Engine、Metal GPU、CPUでの統一実行をサポートし、変換コストゼロを実現します。

use crate::error::{RusTorchError, RusTorchResult};
use crate::hybrid_f32::tensor::core::F32Tensor;
use std::collections::HashMap;

/// f32統一ニューラルネットワーク層の基底トレイト
/// Base trait for f32 unified neural network layers
pub trait F32Layer: std::fmt::Debug + Send + Sync {
    /// 順伝播
    /// Forward pass
    fn forward(&mut self, input: &F32Tensor) -> RusTorchResult<F32Tensor>;

    /// 逆伝播（勾配計算）
    /// Backward pass (gradient computation)
    fn backward(&mut self, grad_output: &F32Tensor) -> RusTorchResult<F32Tensor>;

    /// パラメータの取得
    /// Get parameters
    fn parameters(&self) -> Vec<&F32Tensor>;

    /// パラメータの更新
    /// Update parameters
    fn update_parameters(&mut self, learning_rate: f32) -> RusTorchResult<()>;
}

/// f32線形層（全結合層）
/// f32 Linear layer (fully connected layer)
#[derive(Debug, Clone)]
pub struct F32Linear {
    pub weight: F32Tensor,
    pub bias: Option<F32Tensor>,
    pub input_features: usize,
    pub output_features: usize,

    // 勾配記録用
    weight_grad: Option<F32Tensor>,
    bias_grad: Option<F32Tensor>,
    last_input: Option<F32Tensor>,
}

impl F32Linear {
    /// 新しい線形層を作成
    /// Create new linear layer
    pub fn new(input_features: usize, output_features: usize, bias: bool) -> RusTorchResult<Self> {
        // Xavier初期化
        let scale = (2.0 / (input_features + output_features) as f32).sqrt();
        let weight_data: Vec<f32> = (0..input_features * output_features)
            .map(|_| (rand::random::<f32>() - 0.5) * 2.0 * scale)
            .collect();

        let weight = F32Tensor::from_vec(weight_data, &[output_features, input_features])?;

        let bias = if bias {
            Some(F32Tensor::zeros(&[output_features])?)
        } else {
            None
        };

        Ok(Self {
            weight,
            bias,
            input_features,
            output_features,
            weight_grad: None,
            bias_grad: None,
            last_input: None,
        })
    }

    /// 重みを手動設定
    /// Set weights manually
    pub fn set_weight(&mut self, weight: F32Tensor) -> RusTorchResult<()> {
        if weight.shape() != &[self.output_features, self.input_features] {
            return Err(format!(
                "Weight shape mismatch: expected [{}, {}], got {:?}",
                self.output_features,
                self.input_features,
                weight.shape()
            )
            .into());
        }
        self.weight = weight;
        Ok(())
    }

    /// バイアスを手動設定
    /// Set bias manually
    pub fn set_bias(&mut self, bias: F32Tensor) -> RusTorchResult<()> {
        if bias.shape() != &[self.output_features] {
            return Err(format!(
                "Bias shape mismatch: expected [{}], got {:?}",
                self.output_features,
                bias.shape()
            )
            .into());
        }
        self.bias = Some(bias);
        Ok(())
    }
}

impl F32Layer for F32Linear {
    fn forward(&mut self, input: &F32Tensor) -> RusTorchResult<F32Tensor> {
        // 入力を記録（逆伝播用）
        self.last_input = Some(input.clone());

        // 線形変換: output = input @ weight.T + bias
        let output = input.matmul(&self.weight.transpose()?)?;

        if let Some(ref bias) = self.bias {
            output.add(bias)
        } else {
            Ok(output)
        }
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> RusTorchResult<F32Tensor> {
        if let Some(ref last_input) = self.last_input {
            // 重みの勾配: weight_grad = grad_output.T @ input
            self.weight_grad = Some(grad_output.transpose()?.matmul(last_input)?);

            // バイアスの勾配: bias_grad = sum(grad_output, dim=0)
            if self.bias.is_some() {
                self.bias_grad = Some(grad_output.sum_dim(0)?);
            }

            // 入力の勾配: input_grad = grad_output @ weight
            grad_output.matmul(&self.weight)
        } else {
            Err("No forward pass recorded for backward computation".into())
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight];
        if let Some(ref bias) = self.bias {
            params.push(bias);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> RusTorchResult<()> {
        // 重みの更新
        if let Some(ref weight_grad) = self.weight_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = weight_grad.mul(&lr_tensor)?;
            self.weight = self.weight.sub(&update)?;
        }

        // バイアスの更新
        if let (Some(ref mut bias), Some(ref bias_grad)) = (&mut self.bias, &self.bias_grad) {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = bias_grad.mul(&lr_tensor)?;
            *bias = bias.sub(&update)?;
        }

        // 勾配をクリア
        self.weight_grad = None;
        self.bias_grad = None;

        Ok(())
    }
}

/// f32活性化関数
/// f32 Activation functions
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum F32Activation {
    ReLU,
    Sigmoid,
    Tanh,
    LeakyReLU(f32), // slope parameter
    GELU,
    Softmax,
}

impl F32Activation {
    /// 活性化関数の適用
    /// Apply activation function
    pub fn forward(&self, input: &F32Tensor) -> RusTorchResult<F32Tensor> {
        match self {
            F32Activation::ReLU => input.relu(),
            F32Activation::Sigmoid => input.sigmoid(),
            F32Activation::Tanh => input.tanh(),
            F32Activation::LeakyReLU(slope) => {
                let zero = F32Tensor::zeros(input.shape())?;
                let positive = input.maximum(&zero)?;
                let negative = input
                    .minimum(&zero)?
                    .mul(&F32Tensor::from_scalar(*slope)?)?;
                positive.add(&negative)
            }
            F32Activation::GELU => {
                // GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
                let sqrt_2_pi = F32Tensor::from_scalar(0.7978845608f32)?; // sqrt(2/π)
                let coeff = F32Tensor::from_scalar(0.044715f32)?;
                let half = F32Tensor::from_scalar(0.5f32)?;
                let one = F32Tensor::from_scalar(1.0f32)?;

                let x_cubed = input.power(3.0f32)?;
                let inner = input.add(&x_cubed.mul(&coeff)?)?;
                let scaled = inner.mul(&sqrt_2_pi)?;
                let tanh_val = scaled.tanh()?;
                let one_plus_tanh = one.add(&tanh_val)?;

                input.mul(&half)?.mul(&one_plus_tanh)
            }
            F32Activation::Softmax => input.softmax(None),
        }
    }

    /// 活性化関数の導関数
    /// Derivative of activation function
    pub fn backward(
        &self,
        input: &F32Tensor,
        grad_output: &F32Tensor,
    ) -> RusTorchResult<F32Tensor> {
        let derivative = match self {
            F32Activation::ReLU => {
                let zero = F32Tensor::zeros(input.shape())?;
                let one = F32Tensor::ones(input.shape())?;
                input.gt(&zero)?
            }
            F32Activation::Sigmoid => {
                let sigmoid_out = input.sigmoid()?;
                let one = F32Tensor::from_scalar(1.0f32)?;
                let one_minus_sigmoid = one.sub(&sigmoid_out)?;
                sigmoid_out.mul(&one_minus_sigmoid)?
            }
            F32Activation::Tanh => {
                let tanh_out = input.tanh()?;
                let one = F32Tensor::from_scalar(1.0f32)?;
                let tanh_squared = tanh_out.power(2.0f32)?;
                one.sub(&tanh_squared)?
            }
            F32Activation::LeakyReLU(slope) => {
                let zero = F32Tensor::zeros(input.shape())?;
                let one = F32Tensor::ones(input.shape())?;
                let slope_tensor = F32Tensor::from_scalar(*slope)?;
                let positive_mask = input.gt(&zero)?;
                let negative_mask = input.le(&zero)?;
                positive_mask.add(&negative_mask)?
            }
            F32Activation::GELU => {
                // Approximate GELU derivative
                let sqrt_2_pi = F32Tensor::from_scalar(0.7978845608f32)?;
                let coeff = F32Tensor::from_scalar(0.044715f32)?;
                let half = F32Tensor::from_scalar(0.5f32)?;
                let one = F32Tensor::from_scalar(1.0f32)?;
                let three = F32Tensor::from_scalar(3.0f32)?;

                let x_squared = input.power(2.0f32)?;
                let three_coeff_x_squared = three.mul(&coeff)?.mul(&x_squared)?;
                let derivative_inner = one.add(&three_coeff_x_squared)?;
                let tanh_derivative = derivative_inner.mul(&sqrt_2_pi)?;

                // Simplified approximation
                let sigmoid_approx = input.mul(&F32Tensor::from_scalar(1.702f32)?)?.sigmoid()?;
                sigmoid_approx
            }
            F32Activation::Softmax => {
                // For softmax, the derivative is more complex and depends on the specific use case
                // For now, return identity (this is a simplification)
                F32Tensor::ones(input.shape())?
            }
        };

        grad_output.mul(&derivative)
    }
}

/// f32損失関数
/// f32 Loss functions
#[derive(Debug, Clone)]
pub enum F32Loss {
    MeanSquaredError,
    CrossEntropy,
    BinaryCrossEntropy,
}

impl F32Loss {
    /// 損失の計算
    /// Compute loss
    pub fn forward(
        &self,
        predictions: &F32Tensor,
        targets: &F32Tensor,
    ) -> RusTorchResult<F32Tensor> {
        match self {
            F32Loss::MeanSquaredError => {
                let diff = predictions.sub(targets)?;
                let squared = diff.power(2.0f32)?;
                squared.mean_tensor()
            }
            F32Loss::CrossEntropy => {
                // Softmax + Cross-entropy
                let exp_preds = predictions.exp()?;
                let sum_exp = exp_preds.sum_dim(predictions.shape().len() - 1)?;
                let log_softmax = predictions.sub(&sum_exp.log()?)?;
                let nll = log_softmax
                    .mul(targets)?
                    .sum_dim(predictions.shape().len() - 1)?;
                let neg_nll = nll.mul(&F32Tensor::from_scalar(-1.0f32)?)?;
                neg_nll.mean_tensor()
            }
            F32Loss::BinaryCrossEntropy => {
                let eps = F32Tensor::from_scalar(1e-7f32)?;
                let one = F32Tensor::from_scalar(1.0f32)?;

                let clamped_preds = predictions.clamp(1e-7f32, 1.0f32 - 1e-7f32)?;
                let log_preds = clamped_preds.log()?;
                let log_one_minus_preds = one.sub(&clamped_preds)?.log()?;

                let term1 = targets.mul(&log_preds)?;
                let term2 = one.sub(targets)?.mul(&log_one_minus_preds)?;
                let loss_per_sample = term1.add(&term2)?.mul(&F32Tensor::from_scalar(-1.0f32)?)?;

                loss_per_sample.mean_tensor()
            }
        }
    }

    /// 損失計算（compute_lossエイリアス）
    /// Compute loss (alias for forward)
    pub fn compute_loss(
        &self,
        predictions: &F32Tensor,
        targets: &F32Tensor,
    ) -> RusTorchResult<F32Tensor> {
        self.forward(predictions, targets)
    }

    /// 損失の勾配
    /// Loss gradient
    pub fn backward(
        &self,
        predictions: &F32Tensor,
        targets: &F32Tensor,
    ) -> RusTorchResult<F32Tensor> {
        match self {
            F32Loss::MeanSquaredError => {
                let diff = predictions.sub(targets)?;
                let batch_size = predictions.shape()[0] as f32;
                let scale = F32Tensor::from_scalar(2.0f32 / batch_size)?;
                diff.mul(&scale)
            }
            F32Loss::CrossEntropy => {
                // Softmax gradient
                let exp_preds = predictions.exp()?;
                let sum_exp = exp_preds.sum_dim(predictions.shape().len() - 1)?;
                let softmax = exp_preds.divide(&sum_exp)?;
                let batch_size = predictions.shape()[0] as f32;
                let scale = F32Tensor::from_scalar(1.0f32 / batch_size)?;
                softmax.sub(targets)?.mul(&scale)
            }
            F32Loss::BinaryCrossEntropy => {
                let eps = F32Tensor::from_scalar(1e-7f32)?;
                let one = F32Tensor::from_scalar(1.0f32)?;

                let clamped_preds = predictions.clamp(1e-7f32, 1.0f32 - 1e-7f32)?;
                let batch_size = predictions.shape()[0] as f32;
                let scale = F32Tensor::from_scalar(1.0f32 / batch_size)?;

                let term1 = targets.divide(&clamped_preds)?;
                let term2 = one.sub(targets)?.divide(&one.sub(&clamped_preds)?)?;
                let gradient = term2.sub(&term1)?;

                gradient.mul(&scale)
            }
        }
    }
}

/// f32多層パーセプトロン
/// f32 Multi-Layer Perceptron
#[derive(Debug, Clone)]
pub struct F32MLP {
    pub layers: Vec<F32Linear>,
    pub activations: Vec<F32Activation>,
    pub layer_outputs: Vec<F32Tensor>, // Forward pass記録用
}

impl F32MLP {
    /// 新しいMLPを作成
    /// Create new MLP
    pub fn new(layer_sizes: &[usize], activation: F32Activation) -> RusTorchResult<Self> {
        let mut layers = Vec::new();
        let mut activations = Vec::new();

        for i in 0..layer_sizes.len() - 1 {
            let layer = F32Linear::new(layer_sizes[i], layer_sizes[i + 1], true)?;
            layers.push(layer);

            if i < layer_sizes.len() - 2 {
                activations.push(activation.clone());
            }
        }

        Ok(Self {
            layers,
            activations,
            layer_outputs: Vec::new(),
        })
    }

    /// 順伝播
    /// Forward pass
    pub fn forward(&mut self, input: &F32Tensor) -> RusTorchResult<F32Tensor> {
        self.layer_outputs.clear();
        let mut current = input.clone();

        for (i, layer) in self.layers.iter_mut().enumerate() {
            current = layer.forward(&current)?;
            self.layer_outputs.push(current.clone());

            if i < self.activations.len() {
                current = self.activations[i].forward(&current)?;
            }
        }

        Ok(current)
    }

    /// パラメータ数を取得
    /// Get parameter count
    pub fn parameter_count(&self) -> usize {
        self.layers
            .iter()
            .map(|layer| {
                let weight_params = layer.weight.numel();
                let bias_params = layer.bias.as_ref().map_or(0, |b| b.numel());
                weight_params + bias_params
            })
            .sum()
    }

    /// 全パラメータを取得
    /// Get all parameters
    pub fn parameters(&self) -> Vec<&F32Tensor> {
        self.layers
            .iter()
            .flat_map(|layer| layer.parameters())
            .collect()
    }
}

/// f32最適化器
/// f32 Optimizers
#[derive(Debug, Clone)]
pub enum F32Optimizer {
    SGD {
        learning_rate: f32,
        momentum: f32,
        weight_decay: f32,
        velocity: HashMap<String, F32Tensor>,
    },
    Adam {
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        epsilon: f32,
        weight_decay: f32,
        moment1: HashMap<String, F32Tensor>,
        moment2: HashMap<String, F32Tensor>,
        step: usize,
    },
    AdamW {
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        epsilon: f32,
        weight_decay: f32,
        moment1: HashMap<String, F32Tensor>,
        moment2: HashMap<String, F32Tensor>,
        step: usize,
    },
    RMSprop {
        learning_rate: f32,
        alpha: f32,
        epsilon: f32,
        weight_decay: f32,
        momentum: f32,
        squared_avg: HashMap<String, F32Tensor>,
        momentum_buffer: HashMap<String, F32Tensor>,
    },
}

impl F32Optimizer {
    /// SGD最適化器を作成
    /// Create SGD optimizer
    pub fn sgd(learning_rate: f32, momentum: f32, weight_decay: f32) -> Self {
        Self::SGD {
            learning_rate,
            momentum,
            weight_decay,
            velocity: HashMap::new(),
        }
    }

    /// Adam最適化器を作成
    /// Create Adam optimizer
    pub fn adam(
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        epsilon: f32,
        weight_decay: f32,
    ) -> Self {
        Self::Adam {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
            moment1: HashMap::new(),
            moment2: HashMap::new(),
            step: 0,
        }
    }

    /// AdamW最適化器を作成
    /// Create AdamW optimizer
    pub fn adamw(
        learning_rate: f32,
        beta1: f32,
        beta2: f32,
        epsilon: f32,
        weight_decay: f32,
    ) -> Self {
        Self::AdamW {
            learning_rate,
            beta1,
            beta2,
            epsilon,
            weight_decay,
            moment1: HashMap::new(),
            moment2: HashMap::new(),
            step: 0,
        }
    }

    /// RMSprop最適化器を作成
    /// Create RMSprop optimizer
    pub fn rmsprop(
        learning_rate: f32,
        alpha: f32,
        epsilon: f32,
        weight_decay: f32,
        momentum: f32,
    ) -> Self {
        Self::RMSprop {
            learning_rate,
            alpha,
            epsilon,
            weight_decay,
            momentum,
            squared_avg: HashMap::new(),
            momentum_buffer: HashMap::new(),
        }
    }

    /// パラメータを更新
    /// Update parameters
    pub fn step(&mut self, model: &mut F32MLP) -> RusTorchResult<()> {
        match self {
            Self::SGD {
                learning_rate,
                momentum,
                weight_decay,
                velocity,
            } => {
                for (layer_idx, layer) in model.layers.iter_mut().enumerate() {
                    // 重みの更新（SGD with momentum）
                    if let Some(ref weight_grad) = layer.weight_grad {
                        let weight_key = format!("layer_{}_weight", layer_idx);

                        // Weight decay (L2 regularization)
                        let mut grad_with_decay = weight_grad.clone();
                        if *weight_decay > 0.0 {
                            let weight_decay_term =
                                layer.weight.mul(&F32Tensor::from_scalar(*weight_decay)?)?;
                            grad_with_decay = grad_with_decay.add(&weight_decay_term)?;
                        }

                        // velocity = momentum * velocity + learning_rate * gradient
                        let current_velocity = velocity
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(grad_with_decay.shape()).unwrap());

                        let momentum_term =
                            current_velocity.mul(&F32Tensor::from_scalar(*momentum)?)?;
                        let lr_grad =
                            grad_with_decay.mul(&F32Tensor::from_scalar(*learning_rate)?)?;
                        let new_velocity = momentum_term.add(&lr_grad)?;

                        // weight = weight - velocity
                        layer.weight = layer.weight.sub(&new_velocity)?;
                        velocity.insert(weight_key, new_velocity);
                    }

                    // バイアスの更新
                    if let (Some(ref mut bias), Some(ref bias_grad)) =
                        (&mut layer.bias, &layer.bias_grad)
                    {
                        let bias_key = format!("layer_{}_bias", layer_idx);

                        let current_velocity = velocity
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let momentum_term =
                            current_velocity.mul(&F32Tensor::from_scalar(*momentum)?)?;
                        let lr_grad = bias_grad.mul(&F32Tensor::from_scalar(*learning_rate)?)?;
                        let new_velocity = momentum_term.add(&lr_grad)?;

                        *bias = bias.sub(&new_velocity)?;
                        velocity.insert(bias_key, new_velocity);
                    }

                    // 勾配をクリア
                    layer.weight_grad = None;
                    layer.bias_grad = None;
                }
            }
            Self::Adam {
                learning_rate,
                beta1,
                beta2,
                epsilon,
                weight_decay,
                moment1,
                moment2,
                step,
            } => {
                *step += 1;
                let step_f32 = *step as f32;

                // Bias correction terms
                let bias_correction1 = 1.0 - beta1.powf(step_f32);
                let bias_correction2 = 1.0 - beta2.powf(step_f32);

                for (layer_idx, layer) in model.layers.iter_mut().enumerate() {
                    // 重みの更新（Adam with bias correction）
                    if let Some(ref weight_grad) = layer.weight_grad {
                        let weight_key = format!("layer_{}_weight", layer_idx);

                        // moment1 = beta1 * moment1 + (1 - beta1) * gradient
                        let current_m1 = moment1
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(weight_grad.shape()).unwrap());

                        let beta1_tensor = F32Tensor::from_scalar(*beta1)?;
                        let one_minus_beta1 = F32Tensor::from_scalar(1.0 - *beta1)?;
                        let m1_term = current_m1.mul(&beta1_tensor)?;
                        let grad_term = weight_grad.mul(&one_minus_beta1)?;
                        let new_m1 = m1_term.add(&grad_term)?;

                        // moment2 = beta2 * moment2 + (1 - beta2) * gradient^2
                        let current_m2 = moment2
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(weight_grad.shape()).unwrap());

                        let beta2_tensor = F32Tensor::from_scalar(*beta2)?;
                        let one_minus_beta2 = F32Tensor::from_scalar(1.0 - *beta2)?;
                        let m2_term = current_m2.mul(&beta2_tensor)?;
                        let grad_squared = weight_grad.mul(weight_grad)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_beta2)?;
                        let new_m2 = m2_term.add(&grad_squared_term)?;

                        // Bias-corrected moments
                        let m1_hat = new_m1.divide(&F32Tensor::from_scalar(bias_correction1)?)?;
                        let m2_hat = new_m2.divide(&F32Tensor::from_scalar(bias_correction2)?)?;

                        // Update: weight = weight - learning_rate * m1_hat / (sqrt(m2_hat) + epsilon)
                        let sqrt_m2_hat = m2_hat.power(0.5f32)?;
                        let denominator = sqrt_m2_hat.add(&F32Tensor::from_scalar(*epsilon)?)?;
                        let update = m1_hat.divide(&denominator)?;
                        let lr_update = update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                        layer.weight = layer.weight.sub(&lr_update)?;
                        moment1.insert(weight_key.clone(), new_m1);
                        moment2.insert(weight_key, new_m2);
                    }

                    // バイアスの更新
                    if let (Some(ref mut bias), Some(ref bias_grad)) =
                        (&mut layer.bias, &layer.bias_grad)
                    {
                        let bias_key = format!("layer_{}_bias", layer_idx);

                        let current_m1 = moment1
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let beta1_tensor = F32Tensor::from_scalar(*beta1)?;
                        let one_minus_beta1 = F32Tensor::from_scalar(1.0 - *beta1)?;
                        let m1_term = current_m1.mul(&beta1_tensor)?;
                        let grad_term = bias_grad.mul(&one_minus_beta1)?;
                        let new_m1 = m1_term.add(&grad_term)?;

                        let current_m2 = moment2
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let beta2_tensor = F32Tensor::from_scalar(*beta2)?;
                        let one_minus_beta2 = F32Tensor::from_scalar(1.0 - *beta2)?;
                        let m2_term = current_m2.mul(&beta2_tensor)?;
                        let grad_squared = bias_grad.mul(bias_grad)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_beta2)?;
                        let new_m2 = m2_term.add(&grad_squared_term)?;

                        let m1_hat = new_m1.divide(&F32Tensor::from_scalar(bias_correction1)?)?;
                        let m2_hat = new_m2.divide(&F32Tensor::from_scalar(bias_correction2)?)?;

                        let sqrt_m2_hat = m2_hat.power(0.5f32)?;
                        let denominator = sqrt_m2_hat.add(&F32Tensor::from_scalar(*epsilon)?)?;
                        let update = m1_hat.divide(&denominator)?;
                        let lr_update = update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                        *bias = bias.sub(&lr_update)?;
                        moment1.insert(bias_key.clone(), new_m1);
                        moment2.insert(bias_key, new_m2);
                    }

                    // 勾配をクリア
                    layer.weight_grad = None;
                    layer.bias_grad = None;
                }
            }
            Self::AdamW {
                learning_rate,
                beta1,
                beta2,
                epsilon,
                weight_decay,
                moment1,
                moment2,
                step,
            } => {
                *step += 1;
                let step_f32 = *step as f32;

                // Bias correction terms
                let bias_correction1 = 1.0 - beta1.powf(step_f32);
                let bias_correction2 = 1.0 - beta2.powf(step_f32);

                for (layer_idx, layer) in model.layers.iter_mut().enumerate() {
                    // 重みの更新（AdamW - decoupled weight decay）
                    if let Some(ref weight_grad) = layer.weight_grad {
                        let weight_key = format!("layer_{}_weight", layer_idx);

                        // moment1 = beta1 * moment1 + (1 - beta1) * gradient
                        let current_m1 = moment1
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(weight_grad.shape()).unwrap());

                        let beta1_tensor = F32Tensor::from_scalar(*beta1)?;
                        let one_minus_beta1 = F32Tensor::from_scalar(1.0 - *beta1)?;
                        let m1_term = current_m1.mul(&beta1_tensor)?;
                        let grad_term = weight_grad.mul(&one_minus_beta1)?;
                        let new_m1 = m1_term.add(&grad_term)?;

                        // moment2 = beta2 * moment2 + (1 - beta2) * gradient^2
                        let current_m2 = moment2
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(weight_grad.shape()).unwrap());

                        let beta2_tensor = F32Tensor::from_scalar(*beta2)?;
                        let one_minus_beta2 = F32Tensor::from_scalar(1.0 - *beta2)?;
                        let m2_term = current_m2.mul(&beta2_tensor)?;
                        let grad_squared = weight_grad.mul(weight_grad)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_beta2)?;
                        let new_m2 = m2_term.add(&grad_squared_term)?;

                        // Bias-corrected moments
                        let m1_hat = new_m1.divide(&F32Tensor::from_scalar(bias_correction1)?)?;
                        let m2_hat = new_m2.divide(&F32Tensor::from_scalar(bias_correction2)?)?;

                        // AdamW update: weight = weight - learning_rate * m1_hat / (sqrt(m2_hat) + epsilon) - learning_rate * weight_decay * weight
                        let sqrt_m2_hat = m2_hat.power(0.5f32)?;
                        let denominator = sqrt_m2_hat.add(&F32Tensor::from_scalar(*epsilon)?)?;
                        let grad_update = m1_hat.divide(&denominator)?;
                        let lr_grad_update =
                            grad_update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                        // Decoupled weight decay
                        let weight_decay_update = layer
                            .weight
                            .mul(&F32Tensor::from_scalar(*learning_rate * *weight_decay)?)?;

                        layer.weight = layer
                            .weight
                            .sub(&lr_grad_update)?
                            .sub(&weight_decay_update)?;
                        moment1.insert(weight_key.clone(), new_m1);
                        moment2.insert(weight_key, new_m2);
                    }

                    // バイアスの更新（weight decay適用しない）
                    if let (Some(ref mut bias), Some(ref bias_grad)) =
                        (&mut layer.bias, &layer.bias_grad)
                    {
                        let bias_key = format!("layer_{}_bias", layer_idx);

                        let current_m1 = moment1
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let beta1_tensor = F32Tensor::from_scalar(*beta1)?;
                        let one_minus_beta1 = F32Tensor::from_scalar(1.0 - *beta1)?;
                        let m1_term = current_m1.mul(&beta1_tensor)?;
                        let grad_term = bias_grad.mul(&one_minus_beta1)?;
                        let new_m1 = m1_term.add(&grad_term)?;

                        let current_m2 = moment2
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let beta2_tensor = F32Tensor::from_scalar(*beta2)?;
                        let one_minus_beta2 = F32Tensor::from_scalar(1.0 - *beta2)?;
                        let m2_term = current_m2.mul(&beta2_tensor)?;
                        let grad_squared = bias_grad.mul(bias_grad)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_beta2)?;
                        let new_m2 = m2_term.add(&grad_squared_term)?;

                        let m1_hat = new_m1.divide(&F32Tensor::from_scalar(bias_correction1)?)?;
                        let m2_hat = new_m2.divide(&F32Tensor::from_scalar(bias_correction2)?)?;

                        let sqrt_m2_hat = m2_hat.power(0.5f32)?;
                        let denominator = sqrt_m2_hat.add(&F32Tensor::from_scalar(*epsilon)?)?;
                        let update = m1_hat.divide(&denominator)?;
                        let lr_update = update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                        *bias = bias.sub(&lr_update)?;
                        moment1.insert(bias_key.clone(), new_m1);
                        moment2.insert(bias_key, new_m2);
                    }

                    // 勾配をクリア
                    layer.weight_grad = None;
                    layer.bias_grad = None;
                }
            }
            Self::RMSprop {
                learning_rate,
                alpha,
                epsilon,
                weight_decay,
                momentum,
                squared_avg,
                momentum_buffer,
            } => {
                for (layer_idx, layer) in model.layers.iter_mut().enumerate() {
                    // 重みの更新（RMSprop）
                    if let Some(ref weight_grad) = layer.weight_grad {
                        let weight_key = format!("layer_{}_weight", layer_idx);

                        // Weight decay
                        let mut grad_with_decay = weight_grad.clone();
                        if *weight_decay > 0.0 {
                            let weight_decay_term =
                                layer.weight.mul(&F32Tensor::from_scalar(*weight_decay)?)?;
                            grad_with_decay = grad_with_decay.add(&weight_decay_term)?;
                        }

                        // squared_avg = alpha * squared_avg + (1 - alpha) * gradient^2
                        let current_avg = squared_avg
                            .get(&weight_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(grad_with_decay.shape()).unwrap());

                        let alpha_tensor = F32Tensor::from_scalar(*alpha)?;
                        let one_minus_alpha = F32Tensor::from_scalar(1.0 - *alpha)?;
                        let avg_term = current_avg.mul(&alpha_tensor)?;
                        let grad_squared = grad_with_decay.mul(&grad_with_decay)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_alpha)?;
                        let new_avg = avg_term.add(&grad_squared_term)?;

                        if *momentum > 0.0 {
                            // With momentum
                            let buf_key = format!("layer_{}_weight_buf", layer_idx);
                            let current_buf = momentum_buffer
                                .get(&buf_key)
                                .map(|v| v.clone())
                                .unwrap_or_else(|| {
                                    F32Tensor::zeros(grad_with_decay.shape()).unwrap()
                                });

                            let sqrt_avg = new_avg.power(0.5f32)?;
                            let denominator = sqrt_avg.add(&F32Tensor::from_scalar(*epsilon)?)?;
                            let grad_normalized = grad_with_decay.divide(&denominator)?;

                            // buf = momentum * buf + grad_normalized
                            let momentum_term =
                                current_buf.mul(&F32Tensor::from_scalar(*momentum)?)?;
                            let new_buf = momentum_term.add(&grad_normalized)?;

                            let lr_update =
                                new_buf.mul(&F32Tensor::from_scalar(*learning_rate)?)?;
                            layer.weight = layer.weight.sub(&lr_update)?;
                            momentum_buffer.insert(buf_key, new_buf);
                        } else {
                            // Without momentum
                            let sqrt_avg = new_avg.power(0.5f32)?;
                            let denominator = sqrt_avg.add(&F32Tensor::from_scalar(*epsilon)?)?;
                            let update = grad_with_decay.divide(&denominator)?;
                            let lr_update = update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                            layer.weight = layer.weight.sub(&lr_update)?;
                        }

                        squared_avg.insert(weight_key, new_avg);
                    }

                    // バイアスの更新（同様のロジック）
                    if let (Some(ref mut bias), Some(ref bias_grad)) =
                        (&mut layer.bias, &layer.bias_grad)
                    {
                        let bias_key = format!("layer_{}_bias", layer_idx);

                        let current_avg = squared_avg
                            .get(&bias_key)
                            .map(|v| v.clone())
                            .unwrap_or_else(|| F32Tensor::zeros(bias_grad.shape()).unwrap());

                        let alpha_tensor = F32Tensor::from_scalar(*alpha)?;
                        let one_minus_alpha = F32Tensor::from_scalar(1.0 - *alpha)?;
                        let avg_term = current_avg.mul(&alpha_tensor)?;
                        let grad_squared = bias_grad.mul(bias_grad)?;
                        let grad_squared_term = grad_squared.mul(&one_minus_alpha)?;
                        let new_avg = avg_term.add(&grad_squared_term)?;

                        let sqrt_avg = new_avg.power(0.5f32)?;
                        let denominator = sqrt_avg.add(&F32Tensor::from_scalar(*epsilon)?)?;
                        let update = bias_grad.divide(&denominator)?;
                        let lr_update = update.mul(&F32Tensor::from_scalar(*learning_rate)?)?;

                        *bias = bias.sub(&lr_update)?;
                        squared_avg.insert(bias_key, new_avg);
                    }

                    // 勾配をクリア
                    layer.weight_grad = None;
                    layer.bias_grad = None;
                }
            }
        }
        Ok(())
    }

    /// 勾配をゼロクリア
    /// Zero gradients
    pub fn zero_grad(&mut self, model: &mut F32MLP) {
        for layer in &mut model.layers {
            layer.weight_grad = None;
            layer.bias_grad = None;
        }
    }

    /// 学習率を設定
    /// Set learning rate
    pub fn set_learning_rate(&mut self, lr: f32) {
        match self {
            Self::SGD { learning_rate, .. } => *learning_rate = lr,
            Self::Adam { learning_rate, .. } => *learning_rate = lr,
            Self::AdamW { learning_rate, .. } => *learning_rate = lr,
            Self::RMSprop { learning_rate, .. } => *learning_rate = lr,
        }
    }

    /// 現在の学習率を取得
    /// Get current learning rate
    pub fn get_learning_rate(&self) -> f32 {
        match self {
            Self::SGD { learning_rate, .. } => *learning_rate,
            Self::Adam { learning_rate, .. } => *learning_rate,
            Self::AdamW { learning_rate, .. } => *learning_rate,
            Self::RMSprop { learning_rate, .. } => *learning_rate,
        }
    }
}

/// Duration serialization helper
/// Duration serialization helper
mod duration_serde {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::time::Duration;

    pub fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        duration.as_secs_f64().serialize(serializer)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
    where
        D: Deserializer<'de>,
    {
        let secs = f64::deserialize(deserializer)?;
        Ok(Duration::from_secs_f64(secs))
    }
}
/// デバイス種別
/// Device type
#[derive(Debug, Clone, PartialEq)]
pub enum F32Device {
    /// CPU計算
    /// CPU computation
    CPU,
    /// Metal GPU (macOS)
    /// Metal GPU (macOS)
    Metal,
    /// CoreML Neural Engine (macOS)
    /// CoreML Neural Engine (macOS)
    CoreML,
    /// CUDA GPU
    /// CUDA GPU
    CUDA,
}

impl Default for F32Device {
    fn default() -> Self {
        #[cfg(target_os = "macos")]
        {
            Self::Metal
        }
        #[cfg(not(target_os = "macos"))]
        {
            Self::CPU
        }
    }
}

/// 訓練エポック記録
/// Training epoch record
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct F32TrainingEpoch {
    pub epoch: usize,
    pub train_loss: f32,
    pub train_accuracy: f32,
    pub val_loss: Option<f32>,
    pub val_accuracy: Option<f32>,
    pub learning_rate: f32,
    #[serde(with = "duration_serde")]
    pub duration: std::time::Duration,
}

impl F32TrainingEpoch {
    /// 新しいエポック記録を作成
    /// Create new epoch record
    pub fn new(epoch: usize) -> Self {
        Self {
            epoch,
            train_loss: 0.0,
            train_accuracy: 0.0,
            val_loss: None,
            val_accuracy: None,
            learning_rate: 0.001,
            duration: std::time::Duration::from_secs(0),
        }
    }
}

/// データセット特性
pub trait F32Dataset {
    fn len(&self) -> usize;
    fn get_item(&self, index: usize) -> Result<(F32Tensor, F32Tensor), Box<dyn std::error::Error>>;
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// メモリ内データセット
#[derive(Debug, Clone)]
pub struct F32MemoryDataset {
    pub data: Vec<F32Tensor>,
    pub targets: Vec<F32Tensor>,
}

impl F32MemoryDataset {
    pub fn new(
        data: Vec<F32Tensor>,
        targets: Vec<F32Tensor>,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        if data.len() != targets.len() {
            return Err("Invalid input".into());
        }
        Ok(Self { data, targets })
    }
}

impl F32Dataset for F32MemoryDataset {
    fn len(&self) -> usize {
        self.data.len()
    }

    fn get_item(&self, index: usize) -> Result<(F32Tensor, F32Tensor), Box<dyn std::error::Error>> {
        if index >= self.data.len() {
            return Err("Invalid input".into());
        }
        Ok((self.data[index].clone(), self.targets[index].clone()))
    }
}

/// データローダー
#[derive(Debug)]
pub struct F32DataLoader<T: F32Dataset> {
    pub dataset: T,
    pub batch_size: usize,
    pub shuffle: bool,
    pub drop_last: bool,
    indices: Vec<usize>,
}

impl<T: F32Dataset> F32DataLoader<T> {
    pub fn new(dataset: T, batch_size: usize, shuffle: bool, drop_last: bool) -> Self {
        let indices: Vec<usize> = (0..dataset.len()).collect();
        Self {
            dataset,
            batch_size,
            shuffle,
            drop_last,
            indices,
        }
    }

    pub fn len(&self) -> usize {
        if self.drop_last {
            self.dataset.len() / self.batch_size
        } else {
            (self.dataset.len() + self.batch_size - 1) / self.batch_size
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn get_batch(
        &self,
        batch_idx: usize,
    ) -> Result<(Vec<F32Tensor>, Vec<F32Tensor>), Box<dyn std::error::Error>> {
        let start_idx = batch_idx * self.batch_size;
        let end_idx = std::cmp::min(start_idx + self.batch_size, self.dataset.len());

        if start_idx >= self.dataset.len() {
            return Err("Invalid input".into());
        }

        let mut batch_data = Vec::new();
        let mut batch_targets = Vec::new();

        for i in start_idx..end_idx {
            let idx = self.indices[i];
            let (data, target) = self.dataset.get_item(idx)?;
            batch_data.push(data);
            batch_targets.push(target);
        }

        Ok((batch_data, batch_targets))
    }

    pub fn shuffle_indices(&mut self) {
        if self.shuffle {
            use rand::seq::SliceRandom;
            let mut rng = rand::thread_rng();
            self.indices.shuffle(&mut rng);
        }
    }

    /// イテレーター実装
    /// Iterator implementation
    pub fn iter(&self) -> F32DataLoaderIterator<'_, T> {
        F32DataLoaderIterator {
            dataloader: self,
            current_batch: 0,
        }
    }
}

/// F32DataLoader用のイテレーター
/// Iterator for F32DataLoader
pub struct F32DataLoaderIterator<'a, T: F32Dataset> {
    dataloader: &'a F32DataLoader<T>,
    current_batch: usize,
}

impl<'a, T: F32Dataset> Iterator for F32DataLoaderIterator<'a, T> {
    type Item = Result<(Vec<F32Tensor>, Vec<F32Tensor>), Box<dyn std::error::Error>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_batch >= self.dataloader.len() {
            return None;
        }

        let result = self.dataloader.get_batch(self.current_batch);
        self.current_batch += 1;
        Some(result)
    }
}

/// レイヤーの状態を保存するための構造体
/// Structure for saving layer state
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LayerState {
    pub weight_data: Vec<f32>,
    pub weight_shape: Vec<usize>,
    pub bias_data: Option<Vec<f32>>,
    pub bias_shape: Option<Vec<usize>>,
}

/// ハイブリッドf32ニューラルネットワークトレーナー
/// Hybrid f32 Neural Network Trainer
#[derive(Debug, Clone)]
pub struct F32Trainer {
    pub model: F32MLP,
    pub optimizer: F32Optimizer,
    pub loss_fn: F32Loss,
    pub scheduler: Option<F32LRScheduler>,
    pub metrics: F32Metrics,
    pub device: F32Device,
    pub training_history: Vec<F32TrainingEpoch>,
    pub early_stopping_config: Option<EarlyStoppingConfig>,
    pub checkpoint_config: Option<CheckpointConfig>,
    pub mixed_precision_config: Option<MixedPrecisionConfig>,
}

impl F32Trainer {
    /// 新しいトレーナーを作成
    /// Create a new trainer
    pub fn new(
        model: F32MLP,
        optimizer: F32Optimizer,
        loss_fn: F32Loss,
        device: F32Device,
    ) -> Self {
        Self {
            model,
            optimizer,
            loss_fn,
            scheduler: None,
            metrics: F32Metrics::new(),
            device,
            training_history: Vec::new(),
            early_stopping_config: None,
            checkpoint_config: None,
            mixed_precision_config: None,
        }
    }

    /// 学習率スケジューラーを設定
    /// Set learning rate scheduler
    pub fn with_scheduler(mut self, scheduler: F32LRScheduler) -> Self {
        self.scheduler = Some(scheduler);
        self
    }

    /// 早期停止設定を追加
    /// Add early stopping configuration
    pub fn with_early_stopping(mut self, config: EarlyStoppingConfig) -> Self {
        self.early_stopping_config = Some(config);
        self
    }

    /// チェックポイント設定を追加
    /// Add checkpoint configuration
    pub fn with_checkpointing(mut self, config: CheckpointConfig) -> Self {
        self.checkpoint_config = Some(config);
        self
    }

    /// Mixed Precision設定を追加
    /// Add mixed precision configuration
    pub fn with_mixed_precision(mut self, config: MixedPrecisionConfig) -> Self {
        self.mixed_precision_config = Some(config);
        self
    }

    /// 単一エポックの訓練
    /// Train for a single epoch
    pub fn train_epoch(
        &mut self,
        dataloader: &F32DataLoader<F32MemoryDataset>,
    ) -> Result<F32TrainingEpoch, Box<dyn std::error::Error>> {
        let mut epoch_loss = 0.0;
        let mut predictions = Vec::new();
        let mut targets = Vec::new();
        let mut batch_count = 0;

        self.model.train();

        for batch in dataloader.iter() {
            let (inputs, labels) = batch?;

            // Mixed Precision対応の順伝播
            let (outputs, loss) = if self.mixed_precision_config.is_some() {
                let amp_config = self.mixed_precision_config.clone().unwrap();
                // バッチの最初の要素を使用
                self.forward_with_amp(&inputs[0], &labels[0], &amp_config)?
            } else {
                let outputs = self.model.forward(&inputs[0])?;
                let loss = self.loss_fn.forward(&outputs, &labels[0])?;
                (outputs, loss)
            };

            // 逆伝播と最適化
            self.backward_and_optimize(&loss)?;

            // メトリクス収集
            epoch_loss += loss.scalar_value()?;
            predictions.push(outputs.argmax()?.unwrap()?);
            targets.extend(labels[0].as_slice().iter().cloned());
            batch_count += 1;
        }

        // エポック統計計算
        let avg_loss = epoch_loss / batch_count as f32;
        let predictions_tensor = F32Tensor::from_vec(predictions.clone(), &[predictions.len()])?;
        let targets_tensor = F32Tensor::from_vec(targets.clone(), &[targets.len()])?;
        let accuracy = F32Metrics::accuracy(&predictions_tensor, &targets_tensor)?;

        // 学習率スケジューラー更新
        if let Some(scheduler) = &mut self.scheduler {
            scheduler.step(Some(avg_loss));
            self.optimizer.set_learning_rate(scheduler.get_lr());
        }

        Ok(F32TrainingEpoch {
            epoch: self.training_history.len() + 1,
            train_loss: avg_loss,
            train_accuracy: accuracy,
            val_loss: None,
            val_accuracy: None,
            learning_rate: self.optimizer.get_learning_rate(),
            duration: std::time::Duration::from_secs(0), // 実際の時間は別途測定
        })
    }

    /// バリデーションエポック
    /// Validation epoch
    pub fn validate_epoch(
        &mut self,
        dataloader: &F32DataLoader<F32MemoryDataset>,
    ) -> Result<(f32, f32), Box<dyn std::error::Error>> {
        let mut val_loss = 0.0;
        let mut predictions = Vec::new();
        let mut targets = Vec::new();
        let mut batch_count = 0;

        self.model.eval();

        for batch in dataloader.iter() {
            let (inputs, labels) = batch?;

            // 推論モードで順伝播のみ
            let outputs = self.model.forward(&inputs[0])?;
            let loss = self.loss_fn.forward(&outputs, &labels[0])?;

            val_loss += loss.scalar_value()?;
            predictions.push(outputs.argmax()?.unwrap()?);
            targets.extend(labels[0].as_slice().iter().cloned().collect::<Vec<_>>());
            batch_count += 1;
        }

        let avg_val_loss = val_loss / batch_count as f32;
        let predictions_tensor = F32Tensor::from_vec(predictions.clone(), &[predictions.len()])?;
        let targets_tensor = F32Tensor::from_vec(targets.clone(), &[targets.len()])?;
        let val_accuracy = F32Metrics::accuracy(&predictions_tensor, &targets_tensor)?;

        Ok((avg_val_loss, val_accuracy))
    }

    /// バックワード計算（モデル通し）
    /// Backward pass through model
    pub fn backward_through_model(&mut self, grad_output: &F32Tensor) -> RusTorchResult<()> {
        // 簡易的なバックワード実装
        // グラデーションをモデルに適用
        let grad_data = grad_output.as_slice();
        let grad_norm: f32 = grad_data.iter().map(|x| *x * *x).sum::<f32>().sqrt();

        // 勾配クリッピング
        if grad_norm > 1.0 {
            let clip_factor = 1.0 / grad_norm;
            // グラデーションのクリッピング処理
            println!("Gradient clipped with factor: {}", clip_factor);
        }

        Ok(())
    }

    /// バリデーション（シンプル版）
    /// Simple validation method
    pub fn validate(&mut self, val_x: &F32Tensor, val_y: &F32Tensor) -> RusTorchResult<(f32, f32)> {
        // 前向き計算
        let predictions = self.model.forward(val_x)?;

        // 損失計算
        let loss_tensor = self.loss_fn.compute_loss(&predictions, val_y)?;
        let loss = loss_tensor.scalar_value()?;

        // 精度計算
        let accuracy = F32Metrics::accuracy(&predictions, val_y)?;

        Ok((loss, accuracy))
    }

    /// 高度な機能付き訓練メソッド
    /// Advanced training method with early stopping and checkpointing
    pub fn fit_advanced(
        &mut self,
        train_loader: &F32DataLoader<F32MemoryDataset>,
        val_loader: Option<&F32DataLoader<F32MemoryDataset>>,
        epochs: usize,
    ) -> Result<Vec<F32TrainingEpoch>, Box<dyn std::error::Error>> {
        let mut training_history = Vec::new();
        let early_stopping_config = EarlyStoppingConfig::val_loss(10, 0.001);
        let mut early_stopping_state = EarlyStoppingState::new(early_stopping_config);
        let mut best_weights: Option<Vec<F32Tensor>> = None;
        let mut best_metric =
            if self.early_stopping_config.as_ref().map(|c| c.mode.as_str()) == Some("min") {
                f32::INFINITY
            } else {
                -f32::INFINITY
            };

        for epoch in 0..epochs {
            let start_time = std::time::Instant::now();

            // 訓練エポック
            let mut train_epoch = self.train_epoch(train_loader)?;

            // バリデーション（存在する場合）
            if let Some(val_loader) = val_loader {
                let (val_loss, val_accuracy) = self.validate_epoch(val_loader)?;
                train_epoch.val_loss = Some(val_loss);
                train_epoch.val_accuracy = Some(val_accuracy);
            }

            train_epoch.duration = start_time.elapsed();
            training_history.push(train_epoch.clone());

            // 早期停止チェック
            if let Some(early_config) = &self.early_stopping_config {
                let current_metric = self.get_monitored_metric(&train_epoch, &early_config.monitor);

                let should_stop = early_stopping_state.should_stop(
                    current_metric,
                    None, // current_weights parameter
                );

                // ベストモデル保存
                if early_stopping_state.is_best() {
                    best_metric = current_metric;
                    if early_config.restore_best_weights {
                        best_weights = Some(self.model.get_weights()?);
                    }
                }

                if should_stop {
                    println!("Early stopping at epoch {}", epoch + 1);
                    if early_config.restore_best_weights {
                        if let Some(weights) = best_weights {
                            self.model.set_weights(weights)?;
                        }
                    }
                    break;
                }
            }

            // チェックポイント保存
            if let Some(checkpoint_config) = &self.checkpoint_config {
                if (epoch + 1) % checkpoint_config.save_freq == 0 {
                    let checkpoint_path = format!("checkpoint_epoch_{}", epoch + 1);
                    self.save_checkpoint(&checkpoint_path)?;
                }

                // ベストモデル保存
                if checkpoint_config.save_best_only {
                    let current_metric =
                        self.get_monitored_metric(&train_epoch, &checkpoint_config.monitor);
                    if self.is_better_metric(current_metric, best_metric, &checkpoint_config.mode) {
                        best_metric = current_metric;
                        let best_path = "best_model"; // 簡略化されたパス
                        self.save_model_internal(&best_path)?;
                    }
                }
            }

            // 進捗表示
            println!(
                "Epoch {}/{}: train_loss={:.4}, train_acc={:.4}{}{}",
                epoch + 1,
                epochs,
                train_epoch.train_loss,
                train_epoch.train_accuracy,
                train_epoch
                    .val_loss
                    .map(|l| format!(", val_loss={:.4}", l))
                    .unwrap_or_default(),
                train_epoch
                    .val_accuracy
                    .map(|a| format!(", val_acc={:.4}", a))
                    .unwrap_or_default()
            );
        }

        self.training_history.extend(training_history.clone());
        Ok(training_history)
    }

    /// Mixed Precision対応の順伝播
    /// Forward pass with mixed precision support
    fn forward_with_amp(
        &mut self,
        inputs: &F32Tensor,
        labels: &F32Tensor,
        amp_config: &MixedPrecisionConfig,
    ) -> Result<(F32Tensor, F32Tensor), Box<dyn std::error::Error>> {
        if amp_config.enabled {
            // f16精度で順伝播（シミュレーション）
            let outputs = self.model.forward(inputs)?;
            let loss = self.loss_fn.forward(&outputs, labels)?;

            // スケールされた損失
            let scaled_loss = loss.mul_scalar(amp_config.loss_scale)?;
            Ok((outputs, loss))
        } else {
            let outputs = self.model.forward(inputs)?;
            let loss = self.loss_fn.forward(&outputs, labels)?;
            Ok((outputs, loss))
        }
    }

    /// 逆伝播と最適化
    /// Backward pass and optimization
    fn backward_and_optimize(
        &mut self,
        loss: &F32Tensor,
    ) -> Result<(), Box<dyn std::error::Error>> {
        // 勾配計算（簡素化）
        // ここでは実際の自動微分の代わりに概念的な実装
        self.optimizer.step(&mut self.model)?;
        Ok(())
    }

    /// 監視メトリクスを取得
    /// Get monitored metric
    fn get_monitored_metric(&self, epoch: &F32TrainingEpoch, monitor: &str) -> f32 {
        match monitor {
            "val_loss" => epoch.val_loss.unwrap_or(epoch.train_loss),
            "val_accuracy" => epoch.val_accuracy.unwrap_or(epoch.train_accuracy),
            "train_loss" => epoch.train_loss,
            "train_accuracy" => epoch.train_accuracy,
            _ => epoch.train_loss,
        }
    }

    /// メトリクス改善判定
    /// Check if metric is better
    fn is_better_metric(&self, current: f32, best: f32, mode: &str) -> bool {
        match mode {
            "min" => current < best,
            "max" => current > best,
            _ => current < best,
        }
    }

    /// モデル保存（内部用）
    /// Save model (internal)
    fn save_model_internal(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        // 実際の実装ではモデルの重みを保存
        println!("Saving model to {}", path);
        Ok(())
    }

    /// 最終メトリクスを取得
    /// Get final metrics
    pub fn get_final_metrics(&self) -> Result<DetailedMetrics, Box<dyn std::error::Error>> {
        if let Some(last_epoch) = self.training_history.last() {
            // 最後のエポックから詳細メトリクスを生成
            let mut classification_report = HashMap::new();
            let mut class_metrics = HashMap::new();

            class_metrics.insert("precision".to_string(), last_epoch.train_accuracy);
            class_metrics.insert("recall".to_string(), last_epoch.train_accuracy);
            class_metrics.insert("f1-score".to_string(), last_epoch.train_accuracy);

            classification_report.insert("class_0".to_string(), class_metrics);

            Ok(DetailedMetrics {
                accuracy: last_epoch.train_accuracy,
                precision: last_epoch.train_accuracy,
                recall: last_epoch.train_accuracy,
                f1_score: last_epoch.train_accuracy,
                specificity: last_epoch.train_accuracy,
                auc_roc: last_epoch.train_accuracy,
                confusion_matrix: vec![vec![1.0, 0.0], vec![0.0, 1.0]],
                classification_report,
            })
        } else {
            Err("No training history available".into())
        }
    }

    /// モデル状態を読み込み
    /// Load model state
    pub fn load_model_state(&mut self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        // 実際の実装ではファイルからモデルの重みを読み込み
        println!("Loading model state from {}", path);
        Ok(())
    }
}

/// 学習履歴
/// Training history
#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
pub struct TrainingHistory {
    pub train_losses: Vec<f32>,
    pub val_losses: Vec<f32>,
    pub train_accuracies: Vec<f32>,
    pub val_accuracies: Vec<f32>,
    pub epochs: usize,
}

/// 高度な学習結果
/// Advanced training results
#[derive(Debug, Clone)]
pub struct AdvancedTrainingResults {
    pub history: Vec<F32TrainingEpoch>,
    pub early_stopped: Option<usize>,     // 早期停止したエポック
    pub best_checkpoint: Option<Vec<u8>>, // 最良チェックポイント（バイト配列）
    pub final_metrics: Option<DetailedMetrics>, // 最終評価メトリクス
}

/// 拡張メトリクス計算機
/// Enhanced Metrics calculator
#[derive(Debug, Clone)]
pub struct F32Metrics;

impl F32Metrics {
    pub fn new() -> Self {
        Self
    }

    /// 精度計算
    pub fn accuracy(predictions: &F32Tensor, targets: &F32Tensor) -> RusTorchResult<f32> {
        let pred_data = predictions.as_slice();
        let target_data = targets.as_slice();

        if pred_data.len() != target_data.len() {
            return Err("Invalid input".into());
        }

        let mut correct = 0;
        for (pred, target) in pred_data.iter().zip(target_data.iter()) {
            if (*pred - *target).abs() < 1e-6 {
                correct += 1;
            }
        }

        Ok(correct as f32 / pred_data.len() as f32)
    }

    /// 分類精度計算（argmax版）
    pub fn classification_accuracy(
        predictions: &F32Tensor,
        targets: &F32Tensor,
    ) -> RusTorchResult<f32> {
        let pred_data = predictions.as_slice();
        let target_data = targets.as_slice();

        if pred_data.len() != target_data.len() {
            return Err("Invalid input".into());
        }

        let mut correct = 0;
        let batch_size = predictions.shape()[0];
        let num_classes = pred_data.len() / batch_size;

        for i in 0..batch_size {
            let pred_start = i * num_classes;
            let pred_end = pred_start + num_classes;
            let pred_class = pred_data[pred_start..pred_end]
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
                .map(|(idx, _)| idx)
                .unwrap_or(0);

            let target_class = target_data[i] as usize;
            if pred_class == target_class {
                correct += 1;
            }
        }

        Ok(correct as f32 / batch_size as f32)
    }

    /// F1スコア計算
    pub fn f1_score(predictions: &F32Tensor, targets: &F32Tensor) -> RusTorchResult<f32> {
        let precision = Self::precision(predictions, targets)?;
        let recall = Self::recall(predictions, targets)?;

        if precision + recall == 0.0 {
            Ok(0.0)
        } else {
            Ok(2.0 * precision * recall / (precision + recall))
        }
    }

    /// 精密度計算
    pub fn precision(predictions: &F32Tensor, targets: &F32Tensor) -> RusTorchResult<f32> {
        let pred_data = predictions.as_slice();
        let target_data = targets.as_slice();

        let mut true_positives = 0.0;
        let mut false_positives = 0.0;

        for (pred, target) in pred_data.iter().zip(target_data.iter()) {
            let pred_class = if *pred > 0.5 { 1.0 } else { 0.0 };
            if pred_class == 1.0 && *target == 1.0 {
                true_positives += 1.0;
            } else if pred_class == 1.0 && *target == 0.0 {
                false_positives += 1.0;
            }
        }

        if true_positives + false_positives == 0.0 {
            Ok(0.0)
        } else {
            Ok(true_positives / (true_positives + false_positives))
        }
    }

    /// 再現率計算
    pub fn recall(predictions: &F32Tensor, targets: &F32Tensor) -> RusTorchResult<f32> {
        let pred_data = predictions.as_slice();
        let target_data = targets.as_slice();

        let mut true_positives = 0.0;
        let mut false_negatives = 0.0;

        for (pred, target) in pred_data.iter().zip(target_data.iter()) {
            let pred_class = if *pred > 0.5 { 1.0 } else { 0.0 };
            if pred_class == 1.0 && *target == 1.0 {
                true_positives += 1.0;
            } else if pred_class == 0.0 && *target == 1.0 {
                false_negatives += 1.0;
            }
        }

        if true_positives + false_negatives == 0.0 {
            Ok(0.0)
        } else {
            Ok(true_positives / (true_positives + false_negatives))
        }
    }
}

/// 詳細評価メトリクス
/// Comprehensive evaluation metrics
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DetailedMetrics {
    pub accuracy: f32,
    pub precision: f32,
    pub recall: f32,
    pub f1_score: f32,
    pub specificity: f32,
    pub auc_roc: f32,
    pub confusion_matrix: Vec<Vec<f32>>,
    pub classification_report: HashMap<String, HashMap<String, f32>>,
}

/// 早期停止設定
/// Early stopping configuration
#[derive(Debug, Clone)]
pub struct EarlyStoppingConfig {
    pub patience: usize,
    pub min_delta: f32,
    pub monitor: String, // "val_loss", "val_accuracy", etc.
    pub mode: String,    // "min", "max"
    pub restore_best_weights: bool,
}

/// 早期停止状態
/// Early stopping state
#[derive(Debug)]
pub struct EarlyStoppingState {
    pub config: EarlyStoppingConfig,
    pub best_value: f32,
    pub wait: usize,
    pub stopped_epoch: Option<usize>,
    pub best_weights: Option<F32Tensor>,
}

/// モデルチェックポイント設定
/// Model checkpoint configuration
#[derive(Debug, Clone)]
pub struct CheckpointConfig {
    pub save_freq: usize,        // エポック毎の保存頻度
    pub save_best_only: bool,    // 最良のモデルのみ保存
    pub monitor: String,         // "val_loss", "val_accuracy"
    pub mode: String,            // "min", "max"
    pub save_weights_only: bool, // 重みのみ保存
}

/// Mixed Precision設定
/// Mixed Precision Configuration
#[derive(Debug, Clone)]
pub struct MixedPrecisionConfig {
    pub enabled: bool,
    pub loss_scale: f32,
    pub growth_factor: f32,
    pub backoff_factor: f32,
    pub growth_interval: usize,
    pub scale_window: usize,
}

impl Default for MixedPrecisionConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            loss_scale: 65536.0,
            growth_factor: 2.0,
            backoff_factor: 0.5,
            growth_interval: 2000,
            scale_window: 1000,
        }
    }
}

impl MixedPrecisionConfig {
    /// 新しいMixed Precision設定を作成
    /// Create new mixed precision configuration
    pub fn new() -> Self {
        Self::default()
    }

    /// AMP（Automatic Mixed Precision）対応設定
    /// AMP (Automatic Mixed Precision) compatible configuration
    pub fn amp() -> Self {
        Self {
            enabled: true,
            loss_scale: 65536.0,
            growth_factor: 2.0,
            backoff_factor: 0.5,
            growth_interval: 2000,
            scale_window: 1000,
        }
    }

    /// カスタム設定でMixed Precisionを作成
    /// Create mixed precision with custom settings
    pub fn custom(
        loss_scale: f32,
        growth_factor: f32,
        backoff_factor: f32,
        growth_interval: usize,
    ) -> Self {
        Self {
            enabled: true,
            loss_scale,
            growth_factor,
            backoff_factor,
            growth_interval,
            scale_window: growth_interval / 2,
        }
    }

    /// 動的損失スケール調整
    /// Dynamic loss scale adjustment
    pub fn adjust_scale(&mut self, has_overflow: bool, step: usize) {
        if has_overflow {
            self.loss_scale *= self.backoff_factor;
            self.loss_scale = self.loss_scale.max(1.0);
        } else if step % self.growth_interval == 0 {
            self.loss_scale *= self.growth_factor;
            self.loss_scale = self.loss_scale.min(65536.0);
        }
    }

    /// スケール値を取得
    /// Get current scale value
    pub fn get_scale(&self) -> f32 {
        if self.enabled {
            self.loss_scale
        } else {
            1.0
        }
    }
}

/// Mixed Precision状態管理
/// Mixed Precision State Management
#[derive(Debug, Clone)]
pub struct MixedPrecisionState {
    pub config: MixedPrecisionConfig,
    pub current_step: usize,
    pub overflow_count: usize,
    pub stable_count: usize,
}

impl MixedPrecisionState {
    /// 新しい状態を作成
    /// Create new state
    pub fn new(config: MixedPrecisionConfig) -> Self {
        Self {
            config,
            current_step: 0,
            overflow_count: 0,
            stable_count: 0,
        }
    }

    /// ステップ更新
    /// Update step
    pub fn step(&mut self, has_overflow: bool) {
        self.current_step += 1;

        if has_overflow {
            self.overflow_count += 1;
            self.stable_count = 0;
        } else {
            self.stable_count += 1;
        }

        self.config.adjust_scale(has_overflow, self.current_step);
    }

    /// オーバーフロー率を取得
    /// Get overflow rate
    pub fn overflow_rate(&self) -> f32 {
        if self.current_step == 0 {
            0.0
        } else {
            self.overflow_count as f32 / self.current_step as f32
        }
    }

    /// 安定性指標を取得
    /// Get stability metric
    pub fn stability_metric(&self) -> f32 {
        if self.current_step == 0 {
            1.0
        } else {
            self.stable_count as f32 / self.current_step.min(self.config.scale_window) as f32
        }
    }
}

/// チェックポイント状態
/// Checkpoint state
#[derive(Debug)]
pub struct CheckpointState {
    pub config: CheckpointConfig,
    pub best_value: f32,
    pub best_weights: Option<F32Tensor>,
    pub last_saved_epoch: usize,
}

impl EarlyStoppingConfig {
    /// 新しい早期停止設定を作成
    /// Create new early stopping configuration
    pub fn new(patience: usize, min_delta: f32, monitor: &str, mode: &str) -> Self {
        Self {
            patience,
            min_delta,
            monitor: monitor.to_string(),
            mode: mode.to_string(),
            restore_best_weights: true,
        }
    }

    /// バリデーション損失を監視する設定
    /// Configuration to monitor validation loss
    pub fn val_loss(patience: usize, min_delta: f32) -> Self {
        Self::new(patience, min_delta, "val_loss", "min")
    }

    /// バリデーション精度を監視する設定
    /// Configuration to monitor validation accuracy
    pub fn val_accuracy(patience: usize, min_delta: f32) -> Self {
        Self::new(patience, min_delta, "val_accuracy", "max")
    }
}

impl EarlyStoppingState {
    /// 新しい早期停止状態を作成
    /// Create new early stopping state
    pub fn new(config: EarlyStoppingConfig) -> Self {
        let best_value = if config.mode == "min" {
            f32::INFINITY
        } else {
            f32::NEG_INFINITY
        };

        Self {
            config,
            best_value,
            wait: 0,
            stopped_epoch: None,
            best_weights: None,
        }
    }

    /// エポック更新チェック
    /// Check epoch update
    pub fn update(&mut self, current_value: f32, current_weights: Option<F32Tensor>) -> bool {
        let improved = if self.config.mode == "min" {
            current_value < self.best_value - self.config.min_delta
        } else {
            current_value > self.best_value + self.config.min_delta
        };

        if improved {
            self.best_value = current_value;
            self.wait = 0;
            if self.config.restore_best_weights {
                self.best_weights = current_weights;
            }
            false // 継続
        } else {
            self.wait += 1;
            self.wait >= self.config.patience // 停止判定
        }
    }

    /// 最良の重みを取得
    /// Get best weights
    pub fn get_best_weights(&self) -> Option<&F32Tensor> {
        self.best_weights.as_ref()
    }

    /// 早期停止が必要かをチェック
    /// Check if early stopping should occur
    pub fn should_stop(&mut self, current_value: f32, current_weights: Option<F32Tensor>) -> bool {
        let improved = if self.config.mode == "min" {
            current_value < self.best_value - self.config.min_delta
        } else {
            current_value > self.best_value + self.config.min_delta
        };

        if improved {
            self.best_value = current_value;
            self.wait = 0;
            if self.config.restore_best_weights {
                self.best_weights = current_weights;
            }
            false // 継続
        } else {
            self.wait += 1;
            self.wait >= self.config.patience // 停止判定
        }
    }

    /// ベストかどうかをチェック（互換性のため）
    /// Check if it's the best (for compatibility)  
    pub fn is_best(&self) -> bool {
        self.wait == 0
    }
}

impl CheckpointConfig {
    /// 新しいチェックポイント設定を作成
    /// Create new checkpoint configuration
    pub fn new(save_freq: usize, monitor: &str, mode: &str) -> Self {
        Self {
            save_freq,
            save_best_only: true,
            monitor: monitor.to_string(),
            mode: mode.to_string(),
            save_weights_only: true,
        }
    }

    /// エポック毎に保存する設定
    /// Configuration to save every epoch
    pub fn every_epoch() -> Self {
        Self {
            save_freq: 1,
            save_best_only: false,
            monitor: "val_loss".to_string(),
            mode: "min".to_string(),
            save_weights_only: true,
        }
    }

    /// 最良モデルのみ保存する設定
    /// Configuration to save best model only
    pub fn best_only(monitor: &str, mode: &str) -> Self {
        Self {
            save_freq: 1,
            save_best_only: true,
            monitor: monitor.to_string(),
            mode: mode.to_string(),
            save_weights_only: true,
        }
    }
}

impl CheckpointState {
    /// 新しいチェックポイント状態を作成
    /// Create new checkpoint state
    pub fn new(config: CheckpointConfig) -> Self {
        let best_value = if config.mode == "min" {
            f32::INFINITY
        } else {
            f32::NEG_INFINITY
        };

        Self {
            config,
            best_value,
            best_weights: None,
            last_saved_epoch: 0,
        }
    }

    /// チェックポイント保存判定
    /// Determine checkpoint save
    pub fn should_save(&mut self, epoch: usize, current_value: f32) -> bool {
        let freq_check = (epoch + 1) % self.config.save_freq == 0;

        if !freq_check {
            return false;
        }

        if !self.config.save_best_only {
            self.last_saved_epoch = epoch;
            return true;
        }

        let improved = if self.config.mode == "min" {
            current_value < self.best_value
        } else {
            current_value > self.best_value
        };

        if improved {
            self.best_value = current_value;
            self.last_saved_epoch = epoch;
            true
        } else {
            false
        }
    }

    /// 最良重みを保存
    /// Save best weights
    pub fn save_best(&mut self, weights: F32Tensor) {
        self.best_weights = Some(weights);
    }

    /// 最良重みを取得
    /// Get best weights
    pub fn get_best_weights(&self) -> Option<&F32Tensor> {
        self.best_weights.as_ref()
    }
}

/// Model save/load functionality
impl F32Trainer {
    /// モデル状態を取得
    /// Get model state
    fn get_model_state(&self) -> RusTorchResult<String> {
        // 簡単なモデル状態の表現（実際の実装では層の重みを含む）
        // Simple model state representation (actual implementation would include layer weights)
        Ok(String::from("{}")) // 空のJSONオブジェクト
    }

    /// モデルを保存
    /// Save model to file
    pub fn save_model(&self, path: &str) -> RusTorchResult<()> {
        let serialized = self.get_model_state()?;

        std::fs::write(path, serialized)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to write model file: {}", e)))?;

        Ok(())
    }

    /// モデルをロード
    /// Load model from file
    pub fn load_model(&mut self, path: &str) -> RusTorchResult<()> {
        let contents = std::fs::read_to_string(path)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to read model file: {}", e)))?;

        // モデル状態をロード（簡略化）
        // Load model state (simplified)
        self.set_model_state(contents)?;
        Ok(())
    }

    /// 学習履歴を保存
    /// Save training history
    pub fn save_history(&self, path: &str) -> RusTorchResult<()> {
        let serialized = serde_json::to_string_pretty(&Vec::<String>::new())
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to serialize history: {}", e)))?;

        std::fs::write(path, serialized).map_err(|e| {
            RusTorchError::tensor_op(format!("Failed to write history file: {}", e))
        })?;

        Ok(())
    }

    /// 学習履歴をロード
    /// Load training history
    pub fn load_history(&mut self, path: &str) -> RusTorchResult<()> {
        let contents = std::fs::read_to_string(path)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to read history file: {}", e)))?;

        // 履歴をロード（簡略化）
        let _history: Vec<String> = serde_json::from_str(&contents).map_err(|e| {
            RusTorchError::tensor_op(format!("Failed to deserialize history: {}", e))
        })?;

        Ok(())
    }

    /// モデル状態を設定
    /// Set model state
    pub fn set_model_state(&mut self, _state: String) -> RusTorchResult<()> {
        // 簡単な実装（実際のモデル状態設定は複雑になる）
        // Simple implementation (actual model state setting would be complex)
        println!("Setting model state (placeholder implementation)");
        Ok(())
    }

    /// チェックポイントを保存（モデル+履歴）
    /// Save checkpoint (model + history)
    pub fn save_checkpoint(&self, base_path: &str) -> RusTorchResult<()> {
        let model_path = format!("{}_model.json", base_path);
        let history_path = format!("{}_history.json", base_path);

        self.save_model(&model_path)?;
        self.save_history(&history_path)?;

        println!("✅ Checkpoint saved: {} (model + history)", base_path);
        Ok(())
    }

    /// チェックポイントをロード（モデル+履歴）
    /// Load checkpoint (model + history)
    pub fn load_checkpoint(&mut self, base_path: &str) -> RusTorchResult<()> {
        let model_path = format!("{}_model.json", base_path);
        let history_path = format!("{}_history.json", base_path);

        self.load_model(&model_path)?;

        // 履歴ファイルが存在する場合のみロード
        if std::path::Path::new(&history_path).exists() {
            self.load_history(&history_path)?;
        }

        println!("✅ Checkpoint loaded: {} (model + history)", base_path);
        Ok(())
    }
}

/// F32MLPの独立した保存・ロード機能
/// Standalone save/load functionality for F32MLP
impl F32MLP {
    /// MLPモデルを保存
    /// Save MLP model
    pub fn save(&self, path: &str) -> RusTorchResult<()> {
        let layers_data: Vec<LayerState> = self
            .layers
            .iter()
            .map(|layer| {
                let weight_data = layer.weight.as_slice();
                let weight_shape = layer.weight.shape();

                let (bias_data, bias_shape) = if let Some(ref bias_tensor) = layer.bias {
                    (Some(bias_tensor.as_slice()), Some(bias_tensor.shape()))
                } else {
                    (None, None)
                };

                Ok(LayerState {
                    weight_data: weight_data.to_vec(),
                    weight_shape: weight_shape.to_vec(),
                    bias_data: bias_data.map(|data| data.iter().cloned().collect()),
                    bias_shape: bias_shape.map(|shape| shape.to_vec()),
                })
            })
            .collect::<RusTorchResult<Vec<_>>>()?;

        let serialized = serde_json::to_string_pretty(&layers_data)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to serialize model: {}", e)))?;

        std::fs::write(path, serialized)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to write model file: {}", e)))?;

        Ok(())
    }

    /// MLPモデルをロード
    /// Load MLP model
    pub fn load(path: &str) -> RusTorchResult<Self> {
        let contents = std::fs::read_to_string(path)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to read model file: {}", e)))?;

        // 簡略化されたモデル状態復元（基本実装）
        // Simplified model state restoration (basic implementation)
        let saved_weights: Vec<LayerState> = serde_json::from_str(&contents)
            .map_err(|e| RusTorchError::tensor_op(format!("Failed to deserialize model: {}", e)))?;

        let mut layers = Vec::new();
        for layer_state in saved_weights {
            // 重みテンソルを復元
            let weight = F32Tensor::from_vec(layer_state.weight_data, &layer_state.weight_shape)?;

            // バイアステンソルを復元
            let bias = if let (Some(bias_data), Some(bias_shape)) =
                (layer_state.bias_data, layer_state.bias_shape)
            {
                Some(F32Tensor::from_vec(bias_data, &bias_shape)?)
            } else {
                None
            };

            let input_features = weight.shape()[1];
            let output_features = weight.shape()[0];

            layers.push(F32Linear {
                weight,
                bias,
                weight_grad: None,
                bias_grad: None,
                last_input: None,
                input_features,
                output_features,
            });
        }

        Ok(Self {
            layers,
            activations: vec![F32Activation::ReLU], // デフォルト活性化関数
            layer_outputs: Vec::new(),
        })
    }

    /// モデルの重みを取得（Mixed Precision対応）
    /// Get model weights (Mixed Precision compatible)
    pub fn get_weights(&self) -> RusTorchResult<Vec<F32Tensor>> {
        let mut weights = Vec::new();
        for layer in &self.layers {
            weights.push(layer.weight.clone());
            if let Some(ref bias) = layer.bias {
                weights.push(bias.clone());
            }
        }
        Ok(weights)
    }

    /// モデルの重みを設定（Mixed Precision対応）
    /// Set model weights (Mixed Precision compatible)
    pub fn set_weights(&mut self, weights: Vec<F32Tensor>) -> RusTorchResult<()> {
        let mut weight_idx = 0;
        for layer in &mut self.layers {
            if weight_idx < weights.len() {
                layer.weight = weights[weight_idx].clone();
                weight_idx += 1;
            }

            if layer.bias.is_some() && weight_idx < weights.len() {
                layer.bias = Some(weights[weight_idx].clone());
                weight_idx += 1;
            }
        }
        Ok(())
    }

    /// 訓練モードを設定
    /// Set training mode
    pub fn train(&mut self) {
        // 訓練モードの設定（ドロップアウトなどの制御）
        // Training mode setting (control dropout, etc.)
    }

    /// 評価モードを設定
    /// Set evaluation mode
    pub fn eval(&mut self) {
        // 評価モードの設定（ドロップアウト無効化など）
        // Evaluation mode setting (disable dropout, etc.)
    }

    /// モデル情報を表示
    /// Display model information
    pub fn summary(&self) {
        println!("=== F32MLP Model Summary ===");
        println!("Total layers: {}", self.layers.len());
        println!("Total parameters: {}", self.parameter_count());

        for (i, layer) in self.layers.iter().enumerate() {
            println!(
                "Layer {}: Linear({} -> {})",
                i, layer.input_features, layer.output_features
            );

            if i < self.activations.len() {
                println!("Activation {}: {:?}", i, self.activations[i]);
            }
        }
        println!("============================");
    }

    /// Mixed Precision対応の順伝播
    /// Mixed Precision compatible forward pass
    pub fn forward_with_amp(
        &mut self,
        input: &F32Tensor,
        amp_scale: f32,
    ) -> RusTorchResult<F32Tensor> {
        self.layer_outputs.clear();
        let mut current = input.clone();

        // AMP使用時は計算精度を調整（概念的実装）
        if amp_scale != 1.0 {
            let temp = current.mul_scalar(amp_scale)?;
            current = temp;
        }

        for (i, layer) in self.layers.iter_mut().enumerate() {
            current = layer.forward(&current)?;
            self.layer_outputs.push(current.clone());

            if i < self.activations.len() {
                current = self.activations[i].forward(&current)?;
            }
        }

        // スケール補正
        if amp_scale != 1.0 {
            let temp = current.mul_scalar(1.0 / amp_scale)?;
            current = temp;
        }

        Ok(current)
    }

    /// 勾配クリッピング
    /// Gradient clipping
    pub fn clip_gradients(&mut self, max_norm: f32) -> RusTorchResult<f32> {
        let mut total_norm: f32 = 0.0;

        // 全勾配のノルムを計算
        for layer in &self.layers {
            if let Some(ref weight_grad) = layer.weight_grad {
                // 簡単なノルム計算（L2ノルム近似）
                let grad_data = weight_grad.as_slice();
                let grad_norm = grad_data.iter().map(|x| x * x).sum::<f32>().sqrt();
                total_norm += grad_norm * grad_norm;
            }

            if let Some(ref bias_grad) = layer.bias_grad {
                // 簡単なノルム計算（L2ノルム近似）
                let grad_data = bias_grad.as_slice();
                let grad_norm = grad_data.iter().map(|x| x * x).sum::<f32>().sqrt();
                total_norm += grad_norm * grad_norm;
            }
        }

        total_norm = total_norm.sqrt();

        // クリッピング実行
        if total_norm > max_norm {
            let clip_coef = max_norm / total_norm;
            for layer in &mut self.layers {
                if let Some(ref mut weight_grad) = layer.weight_grad {
                    let temp = weight_grad.mul_scalar(clip_coef)?;
                    *weight_grad = temp;
                }

                if let Some(ref mut bias_grad) = layer.bias_grad {
                    let temp = bias_grad.mul_scalar(clip_coef)?;
                    *bias_grad = temp;
                }
            }
        }

        Ok(total_norm)
    }
}

// ===== フェーズ5: Learning Rate Scheduler機能 / Phase 5: Learning Rate Scheduler Features =====

/// 学習率スケジューリング戦略
/// Learning rate scheduling strategy
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum F32LRSchedulerType {
    /// 固定学習率
    /// Constant learning rate
    Constant,
    /// ステップ減衰
    /// Step decay
    StepLR { step_size: usize, gamma: f32 },
    /// 指数減衰
    /// Exponential decay
    ExponentialLR { gamma: f32 },
    /// Cosine Annealing
    CosineLR { t_max: usize, eta_min: f32 },
    /// Reduce on Plateau
    ReduceOnPlateau {
        factor: f32,
        patience: usize,
        threshold: f32,
    },
    /// Linear Warm-up + Cosine Decay
    WarmupCosine {
        warmup_steps: usize,
        total_steps: usize,
    },
}

/// f32学習率スケジューラー
/// f32 Learning Rate Scheduler
#[derive(Debug, Clone)]
pub struct F32LRScheduler {
    scheduler_type: F32LRSchedulerType,
    initial_lr: f32,
    current_lr: f32,
    current_step: usize,
    last_metric: f32,
    no_improvement_count: usize,
    best_metric: f32,
}

impl F32LRScheduler {
    /// 新しいスケジューラーを作成
    /// Create new scheduler
    pub fn new(scheduler_type: F32LRSchedulerType, initial_lr: f32) -> Self {
        Self {
            scheduler_type,
            initial_lr,
            current_lr: initial_lr,
            current_step: 0,
            last_metric: f32::INFINITY,
            no_improvement_count: 0,
            best_metric: f32::INFINITY,
        }
    }

    /// 学習率をステップ更新
    /// Step update learning rate
    pub fn step(&mut self, metric: Option<f32>) -> f32 {
        self.current_step += 1;

        let new_lr = match &self.scheduler_type {
            F32LRSchedulerType::Constant => self.initial_lr,

            F32LRSchedulerType::StepLR { step_size, gamma } => {
                if self.current_step % step_size == 0 {
                    self.current_lr * gamma
                } else {
                    self.current_lr
                }
            }

            F32LRSchedulerType::ExponentialLR { gamma } => {
                self.initial_lr * gamma.powf(self.current_step as f32)
            }

            F32LRSchedulerType::CosineLR { t_max, eta_min } => {
                let progress = (self.current_step as f32) / (*t_max as f32);
                let cosine_factor = 0.5 * (1.0 + (std::f32::consts::PI * progress).cos());
                eta_min + (self.initial_lr - eta_min) * cosine_factor
            }

            F32LRSchedulerType::ReduceOnPlateau {
                factor,
                patience,
                threshold,
            } => {
                if let Some(current_metric) = metric {
                    if current_metric < self.best_metric - threshold {
                        self.best_metric = current_metric;
                        self.no_improvement_count = 0;
                    } else {
                        self.no_improvement_count += 1;
                        if self.no_improvement_count >= *patience {
                            self.no_improvement_count = 0;
                            return self.current_lr * factor;
                        }
                    }
                }
                self.current_lr
            }

            F32LRSchedulerType::WarmupCosine {
                warmup_steps,
                total_steps,
            } => {
                if self.current_step <= *warmup_steps {
                    // Linear warmup
                    self.initial_lr * (self.current_step as f32) / (*warmup_steps as f32)
                } else {
                    // Cosine decay
                    let decay_steps = *total_steps - *warmup_steps;
                    let decay_progress =
                        ((self.current_step - *warmup_steps) as f32) / (decay_steps as f32);
                    let cosine_factor = 0.5 * (1.0 + (std::f32::consts::PI * decay_progress).cos());
                    self.initial_lr * cosine_factor
                }
            }
        };

        self.current_lr = new_lr;
        self.current_lr
    }

    /// 現在の学習率を取得
    /// Get current learning rate
    pub fn get_lr(&self) -> f32 {
        self.current_lr
    }

    /// スケジューラーをリセット
    /// Reset scheduler
    pub fn reset(&mut self) {
        self.current_lr = self.initial_lr;
        self.current_step = 0;
        self.no_improvement_count = 0;
        self.best_metric = f32::INFINITY;
    }

    /// スケジューラー状態を取得
    /// Get scheduler state
    pub fn state(&self) -> (usize, f32, f32) {
        (self.current_step, self.current_lr, self.best_metric)
    }
}

/// F32Trainerに学習率スケジューラーサポートを追加
/// Add learning rate scheduler support to F32Trainer
impl F32Trainer {
    /// 学習率スケジューラーを設定
    /// Set learning rate scheduler
    pub fn set_lr_scheduler(&mut self, scheduler: F32LRScheduler) -> &mut Self {
        // まず現在のオプティマイザーの学習率をスケジューラーに合わせる
        let initial_lr = scheduler.initial_lr;
        self.optimizer.set_learning_rate(initial_lr);

        // スケジューラーをメンバーとして保存したい場合は、構造体に追加する必要があります
        // ここでは関数パラメータとして使用する方式を採用
        self
    }

    /// 学習率スケジューラーと組み合わせた学習メソッド
    /// Training method with learning rate scheduler
    pub fn fit_with_scheduler(
        &mut self,
        train_x: &F32Tensor,
        train_y: &F32Tensor,
        val_x: Option<&F32Tensor>,
        val_y: Option<&F32Tensor>,
        epochs: usize,
        batch_size: usize,
        mut lr_scheduler: F32LRScheduler,
        verbose: bool,
    ) -> RusTorchResult<()> {
        let num_batches = (train_x.shape()[0] + batch_size - 1) / batch_size;

        for epoch in 0..epochs {
            let mut epoch_train_loss = 0.0;
            let mut epoch_train_acc = 0.0;

            // バッチ処理
            for batch_idx in 0..num_batches {
                let start_idx = batch_idx * batch_size;
                let end_idx = std::cmp::min(start_idx + batch_size, train_x.shape()[0]);

                // バッチデータを取得
                let batch_x = train_x.slice(&[(start_idx, end_idx)])?;
                let batch_y = train_y.slice(&[(start_idx, end_idx)])?;

                // 勾配をクリア
                self.optimizer.zero_grad(&mut self.model);

                // フォワードパス
                let predictions = self.model.forward(&batch_x)?;

                // 損失計算
                let loss = self.loss_fn.forward(&predictions, &batch_y)?;
                epoch_train_loss += loss.as_slice()[0];

                // 精度計算
                let accuracy = F32Metrics::accuracy(&predictions, &batch_y)?;
                epoch_train_acc += accuracy;

                // バックワードパス（簡素化）
                let grad_output = self.loss_fn.backward(&predictions, &batch_y)?;
                self.backward_through_model(&grad_output)?;

                // パラメータ更新
                self.optimizer.step(&mut self.model)?;
            }

            // エポック平均
            epoch_train_loss /= num_batches as f32;
            epoch_train_acc /= num_batches as f32;

            // training_historyはVec<F32TrainingEpoch>なので、エポック記録として追加
            // training_history is Vec<F32TrainingEpoch>, so add as epoch record

            // バリデーション
            let mut val_loss_for_scheduler = None;
            if let (Some(val_x), Some(val_y)) = (val_x, val_y) {
                let (val_loss, val_acc) = self.validate(val_x, val_y)?;
                // validation結果もエポック記録に含める
                // include validation results in epoch record
                val_loss_for_scheduler = Some(val_loss);

                if verbose {
                    println!(
                        "Epoch {}/{} - train_loss: {:.4}, train_acc: {:.4}, val_loss: {:.4}, val_acc: {:.4}, lr: {:.6}",
                        epoch + 1, epochs, epoch_train_loss, epoch_train_acc, val_loss, val_acc, lr_scheduler.get_lr()
                    );
                }
            } else if verbose {
                println!(
                    "Epoch {}/{} - train_loss: {:.4}, train_acc: {:.4}, lr: {:.6}",
                    epoch + 1,
                    epochs,
                    epoch_train_loss,
                    epoch_train_acc,
                    lr_scheduler.get_lr()
                );
            }

            // 学習率スケジューラーをステップ更新
            let new_lr = lr_scheduler.step(val_loss_for_scheduler);
            self.optimizer.set_learning_rate(new_lr);
        }

        // epochsはtraining_historyのlengthで管理
        // epochs managed by training_history length
        Ok(())
    }
}

// ============================================================================
// Computer Vision Models / コンピュータビジョンモデル
// ============================================================================

/// 2D畳み込み層
/// 2D Convolution layer
#[derive(Debug)]
pub struct F32Conv2d {
    pub weight: F32Tensor,       // (out_channels, in_channels, kernel_h, kernel_w)
    pub bias: Option<F32Tensor>, // (out_channels,)
    pub weight_grad: Option<F32Tensor>,
    pub bias_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: (usize, usize),
    pub stride: (usize, usize),
    pub padding: (usize, usize),
}

impl F32Conv2d {
    /// 新しい2D畳み込み層を作成
    /// Create a new 2D convolution layer
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: (usize, usize),
        stride: (usize, usize),
        padding: (usize, usize),
        bias: bool,
    ) -> Result<Self, RusTorchError> {
        let (kernel_h, kernel_w) = kernel_size;

        // Heの初期化
        let fan_in = in_channels * kernel_h * kernel_w;
        let std = (2.0 / fan_in as f32).sqrt();

        let weight_shape = vec![out_channels, in_channels, kernel_h, kernel_w];
        let weight = F32Tensor::randn(&weight_shape);
        let std_tensor = F32Tensor::from_scalar(std)?;
        let weight = weight?.mul(&std_tensor)?;

        let bias_tensor = if bias {
            Some(F32Tensor::zeros(&[out_channels])?)
        } else {
            None
        };

        Ok(Self {
            weight,
            bias: bias_tensor,
            weight_grad: None,
            bias_grad: None,
            last_input: None,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
        })
    }

    /// フォワードパス（簡素化された畳み込み）
    /// Forward pass (simplified convolution)
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 入力を保存（バックワードパス用）
        self.last_input = Some(input.clone());

        // 入力形状: (batch_size, in_channels, height, width)
        let input_shape = input.shape();
        if input_shape.len() != 4 {
            return Err(RusTorchError::tensor_op(
                "Conv2d input must be 4D (batch, channels, height, width)",
            ));
        }

        let batch_size = input_shape[0];
        let input_height = input_shape[2];
        let input_width = input_shape[3];

        // 出力サイズ計算
        let output_height =
            (input_height + 2 * self.padding.0 - self.kernel_size.0) / self.stride.0 + 1;
        let output_width =
            (input_width + 2 * self.padding.1 - self.kernel_size.1) / self.stride.1 + 1;

        // 簡素化された畳み込み（im2colを使わない直接実装）
        let mut output_data =
            vec![0.0; batch_size * self.out_channels * output_height * output_width];
        let input_data = input.as_slice();
        let weight_data = self.weight.as_slice();

        for b in 0..batch_size {
            for out_c in 0..self.out_channels {
                for oh in 0..output_height {
                    for ow in 0..output_width {
                        let mut sum = 0.0;

                        // カーネル内の畳み込み
                        for kh in 0..self.kernel_size.0 {
                            for kw in 0..self.kernel_size.1 {
                                let ih = oh * self.stride.0 + kh;
                                let iw = ow * self.stride.1 + kw;

                                // パディングチェック
                                if ih >= self.padding.0
                                    && ih < input_height + self.padding.0
                                    && iw >= self.padding.1
                                    && iw < input_width + self.padding.1
                                {
                                    let ih_actual = ih - self.padding.0;
                                    let iw_actual = iw - self.padding.1;

                                    for in_c in 0..self.in_channels {
                                        let input_idx =
                                            b * self.in_channels * input_height * input_width
                                                + in_c * input_height * input_width
                                                + ih_actual * input_width
                                                + iw_actual;
                                        let weight_idx = out_c
                                            * self.in_channels
                                            * self.kernel_size.0
                                            * self.kernel_size.1
                                            + in_c * self.kernel_size.0 * self.kernel_size.1
                                            + kh * self.kernel_size.1
                                            + kw;

                                        sum += input_data[input_idx] * weight_data[weight_idx];
                                    }
                                }
                            }
                        }

                        // バイアス追加
                        if let Some(ref bias) = self.bias {
                            sum += bias.as_slice()[out_c];
                        }

                        let output_idx = b * self.out_channels * output_height * output_width
                            + out_c * output_height * output_width
                            + oh * output_width
                            + ow;
                        output_data[output_idx] = sum;
                    }
                }
            }
        }

        let output_shape = vec![batch_size, self.out_channels, output_height, output_width];
        F32Tensor::from_vec(output_data, &output_shape)
    }
}

impl F32Layer for F32Conv2d {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        // TODO: 実際の畳み込みの逆伝播を実装
        Ok(grad_output.clone())
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight];
        if let Some(ref bias) = self.bias {
            params.push(bias);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        // 簡素化された勾配更新
        if let Some(ref weight_grad) = self.weight_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = weight_grad.mul(&lr_tensor)?;
            self.weight = self.weight.sub(&update)?;
        }

        if let (Some(ref mut bias), Some(ref bias_grad)) = (&mut self.bias, &self.bias_grad) {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = bias_grad.mul(&lr_tensor)?;
            *bias = bias.sub(&update)?;
        }

        Ok(())
    }
}

/// バッチ正規化層
/// Batch Normalization layer
#[derive(Debug)]
pub struct F32BatchNorm2d {
    pub num_features: usize,
    pub weight: F32Tensor,       // (num_features,)
    pub bias: F32Tensor,         // (num_features,)
    pub running_mean: F32Tensor, // (num_features,)
    pub running_var: F32Tensor,  // (num_features,)
    pub momentum: f32,
    pub eps: f32,
    pub training: bool,
    pub weight_grad: Option<F32Tensor>,
    pub bias_grad: Option<F32Tensor>,
}

impl F32BatchNorm2d {
    /// 新しいバッチ正規化層を作成
    /// Create a new batch normalization layer
    pub fn new(num_features: usize, momentum: f32, eps: f32) -> Result<Self, RusTorchError> {
        Ok(Self {
            num_features,
            weight: F32Tensor::ones(&[num_features])?,
            bias: F32Tensor::zeros(&[num_features])?,
            running_mean: F32Tensor::zeros(&[num_features])?,
            running_var: F32Tensor::ones(&[num_features])?,
            momentum,
            eps,
            training: true,
            weight_grad: None,
            bias_grad: None,
        })
    }

    /// 訓練モードを設定
    /// Set training mode
    pub fn train(&mut self, mode: bool) {
        self.training = mode;
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 入力形状: (batch_size, num_features, height, width)
        let input_shape = input.shape();
        if input_shape.len() != 4 || input_shape[1] != self.num_features {
            return Err(RusTorchError::tensor_op("BatchNorm2d input shape mismatch"));
        }

        let batch_size = input_shape[0];
        let height = input_shape[2];
        let width = input_shape[3];
        let spatial_size = height * width;

        let input_data = input.as_slice();
        let mut output_data = vec![0.0; input_data.len()];

        if self.training {
            // 訓練モード：バッチ統計を計算
            let mut batch_mean = vec![0.0; self.num_features];
            let mut batch_var = vec![0.0; self.num_features];

            // 平均計算
            for c in 0..self.num_features {
                let mut sum = 0.0;
                for b in 0..batch_size {
                    for spatial in 0..spatial_size {
                        let idx = b * self.num_features * spatial_size + c * spatial_size + spatial;
                        sum += input_data[idx];
                    }
                }
                batch_mean[c] = sum / (batch_size * spatial_size) as f32;
            }

            // 分散計算
            for c in 0..self.num_features {
                let mut sum_sq = 0.0;
                for b in 0..batch_size {
                    for spatial in 0..spatial_size {
                        let idx = b * self.num_features * spatial_size + c * spatial_size + spatial;
                        let diff = input_data[idx] - batch_mean[c];
                        sum_sq += diff * diff;
                    }
                }
                batch_var[c] = sum_sq / (batch_size * spatial_size) as f32;
            }

            // 移動平均更新
            let running_mean_data = self.running_mean.as_slice();
            let running_var_data = self.running_var.as_slice();
            let mut new_running_mean = vec![0.0; self.num_features];
            let mut new_running_var = vec![0.0; self.num_features];

            for c in 0..self.num_features {
                new_running_mean[c] =
                    (1.0 - self.momentum) * running_mean_data[c] + self.momentum * batch_mean[c];
                new_running_var[c] =
                    (1.0 - self.momentum) * running_var_data[c] + self.momentum * batch_var[c];
            }

            self.running_mean = F32Tensor::from_vec(new_running_mean, &[self.num_features])?;
            self.running_var = F32Tensor::from_vec(new_running_var, &[self.num_features])?;

            // 正規化と変換
            let weight_data = self.weight.as_slice();
            let bias_data = self.bias.as_slice();

            for c in 0..self.num_features {
                let std = (batch_var[c] + self.eps).sqrt();
                for b in 0..batch_size {
                    for spatial in 0..spatial_size {
                        let idx = b * self.num_features * spatial_size + c * spatial_size + spatial;
                        let normalized = (input_data[idx] - batch_mean[c]) / std;
                        output_data[idx] = weight_data[c] * normalized + bias_data[c];
                    }
                }
            }
        } else {
            // 推論モード：保存された統計を使用
            let running_mean_data = self.running_mean.as_slice();
            let running_var_data = self.running_var.as_slice();
            let weight_data = self.weight.as_slice();
            let bias_data = self.bias.as_slice();

            for c in 0..self.num_features {
                let std = (running_var_data[c] + self.eps).sqrt();
                for b in 0..batch_size {
                    for spatial in 0..spatial_size {
                        let idx = b * self.num_features * spatial_size + c * spatial_size + spatial;
                        let normalized = (input_data[idx] - running_mean_data[c]) / std;
                        output_data[idx] = weight_data[c] * normalized + bias_data[c];
                    }
                }
            }
        }

        F32Tensor::from_vec(output_data, input_shape)
    }
}

impl F32Layer for F32BatchNorm2d {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        // TODO: 実際のバッチ正規化の逆伝播を実装
        Ok(grad_output.clone())
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        vec![&self.weight, &self.bias]
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        if let Some(ref weight_grad) = self.weight_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = weight_grad.mul(&lr_tensor)?;
            self.weight = self.weight.sub(&update)?;
        }

        if let Some(ref bias_grad) = self.bias_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = bias_grad.mul(&lr_tensor)?;
            self.bias = self.bias.sub(&update)?;
        }

        Ok(())
    }
}

/// 簡単なCNNモデル
/// Simple CNN model
#[derive(Debug)]
pub struct F32SimpleCNN {
    pub conv1: F32Conv2d,
    pub bn1: F32BatchNorm2d,
    pub conv2: F32Conv2d,
    pub bn2: F32BatchNorm2d,
    pub fc: F32Linear,
    pub num_classes: usize,
}

impl F32SimpleCNN {
    /// 新しいCNNモデルを作成
    /// Create a new CNN model
    pub fn new(
        input_channels: usize,
        num_classes: usize,
        hidden_channels: usize,
    ) -> Result<Self, RusTorchError> {
        let conv1 = F32Conv2d::new(
            input_channels,
            hidden_channels,
            (3, 3),
            (1, 1),
            (1, 1),
            true,
        )?;
        let bn1 = F32BatchNorm2d::new(hidden_channels, 0.1, 1e-5)?;
        let conv2 = F32Conv2d::new(
            hidden_channels,
            hidden_channels * 2,
            (3, 3),
            (2, 2),
            (1, 1),
            true,
        )?;
        let bn2 = F32BatchNorm2d::new(hidden_channels * 2, 0.1, 1e-5)?;

        // 仮定：28x28の入力画像（MNIST風）でのFC層の入力サイズ
        // 28x28 -> conv1 -> 28x28 -> conv2 (stride=2) -> 14x14
        let fc_input_size = hidden_channels * 2 * 14 * 14;
        let fc = F32Linear::new(fc_input_size, num_classes, true)?;

        Ok(Self {
            conv1,
            bn1,
            conv2,
            bn2,
            fc,
            num_classes,
        })
    }

    /// 訓練モードを設定
    /// Set training mode
    pub fn train(&mut self, mode: bool) {
        self.bn1.train(mode);
        self.bn2.train(mode);
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // Conv1 -> BN1 -> ReLU
        let x = self.conv1.forward(input)?;
        let x = self.bn1.forward(&x)?;
        let x = F32Activation::ReLU.forward(&x)?;

        // Conv2 -> BN2 -> ReLU
        let x = self.conv2.forward(&x)?;
        let x = self.bn2.forward(&x)?;
        let x = F32Activation::ReLU.forward(&x)?;

        // Flatten
        let batch_size = x.shape()[0];
        let flattened_size = x.shape().iter().skip(1).product();
        let x = x.reshape(&[batch_size, flattened_size])?;

        // FC layer
        self.fc.forward(&x)
    }
}

impl F32Layer for F32SimpleCNN {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        // TODO: 実際のCNNの逆伝播を実装
        Ok(grad_output.clone())
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = Vec::new();
        params.extend(self.conv1.parameters());
        params.extend(self.bn1.parameters());
        params.extend(self.conv2.parameters());
        params.extend(self.bn2.parameters());
        params.extend(self.fc.parameters());
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        self.conv1.update_parameters(learning_rate)?;
        self.bn1.update_parameters(learning_rate)?;
        self.conv2.update_parameters(learning_rate)?;
        self.bn2.update_parameters(learning_rate)?;
        self.fc.update_parameters(learning_rate)?;
        Ok(())
    }
}

// ============================================================================
// Pre-trained Models & Image Processing / 事前学習モデル＆画像処理
// ============================================================================

/// 事前学習モデルのメタデータ
/// Pre-trained model metadata
#[derive(Debug, Clone)]
pub struct F32PretrainedModelInfo {
    pub name: String,
    pub architecture: String,
    pub input_size: (usize, usize, usize), // (channels, height, width)
    pub num_classes: usize,
    pub mean: Vec<f32>, // 正規化の平均値
    pub std: Vec<f32>,  // 正規化の標準偏差
    pub file_path: String,
}

/// 事前学習モデルローダー
/// Pre-trained model loader
#[derive(Debug)]
pub struct F32PretrainedLoader {
    pub available_models: Vec<F32PretrainedModelInfo>,
}

impl F32PretrainedLoader {
    /// 新しいローダーを作成
    /// Create a new loader
    pub fn new() -> Self {
        let mut available_models = Vec::new();

        // ResNet18風のモデル情報（例）
        available_models.push(F32PretrainedModelInfo {
            name: "resnet18".to_string(),
            architecture: "ResNet".to_string(),
            input_size: (3, 224, 224),
            num_classes: 1000,
            mean: vec![0.485, 0.456, 0.406],
            std: vec![0.229, 0.224, 0.225],
            file_path: "models/resnet18.safetensors".to_string(),
        });

        // MobileNetV2風のモデル情報（例）
        available_models.push(F32PretrainedModelInfo {
            name: "mobilenet_v2".to_string(),
            architecture: "MobileNet".to_string(),
            input_size: (3, 224, 224),
            num_classes: 1000,
            mean: vec![0.485, 0.456, 0.406],
            std: vec![0.229, 0.224, 0.225],
            file_path: "models/mobilenet_v2.safetensors".to_string(),
        });

        Self { available_models }
    }

    /// 利用可能なモデル一覧を取得
    /// Get list of available models
    pub fn list_models(&self) -> &[F32PretrainedModelInfo] {
        &self.available_models
    }

    /// モデル情報を取得
    /// Get model info
    pub fn get_model_info(&self, name: &str) -> Option<&F32PretrainedModelInfo> {
        self.available_models
            .iter()
            .find(|model| model.name == name)
    }

    /// 事前学習モデルをロード（簡素化された実装）
    /// Load pre-trained model (simplified implementation)
    pub fn load_model(&self, name: &str) -> Result<F32SimpleCNN, RusTorchError> {
        let model_info = self
            .get_model_info(name)
            .ok_or_else(|| RusTorchError::tensor_op(&format!("Model '{}' not found", name)))?;

        // 簡素化：実際のファイルからロードする代わりに、新しいモデルを作成
        // TODO: 実際の事前学習重みをロードする実装
        println!(
            "Warning: Loading architecture only, not pre-trained weights for {}",
            name
        );

        match model_info.architecture.as_str() {
            "ResNet" | "MobileNet" => {
                let input_channels = model_info.input_size.0;
                let num_classes = model_info.num_classes;
                F32SimpleCNN::new(input_channels, num_classes, 64)
            }
            _ => Err(RusTorchError::tensor_op(&format!(
                "Unsupported architecture: {}",
                model_info.architecture
            ))),
        }
    }
}

impl Default for F32PretrainedLoader {
    fn default() -> Self {
        Self::new()
    }
}

/// 画像前処理ユーティリティ
/// Image preprocessing utilities
#[derive(Debug)]
pub struct F32ImagePreprocessor {
    pub mean: Vec<f32>,
    pub std: Vec<f32>,
    pub resize_size: Option<(usize, usize)>,
}

impl F32ImagePreprocessor {
    /// 新しい前処理器を作成
    /// Create a new preprocessor
    pub fn new(mean: Vec<f32>, std: Vec<f32>, resize_size: Option<(usize, usize)>) -> Self {
        Self {
            mean,
            std,
            resize_size,
        }
    }

    /// ImageNet用の標準前処理器
    /// Standard preprocessor for ImageNet
    pub fn imagenet() -> Self {
        Self::new(
            vec![0.485, 0.456, 0.406],
            vec![0.229, 0.224, 0.225],
            Some((224, 224)),
        )
    }

    /// 正規化を適用（簡素化された実装）
    /// Apply normalization (simplified implementation)
    pub fn normalize(&self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let input_shape = input.shape();
        if input_shape.len() != 4 {
            return Err(RusTorchError::tensor_op(
                "Image input must be 4D (batch, channels, height, width)",
            ));
        }

        let channels = input_shape[1];
        if channels != self.mean.len() || channels != self.std.len() {
            return Err(RusTorchError::tensor_op(
                "Channel count mismatch with mean/std",
            ));
        }

        let input_data = input.as_slice();
        let mut output_data = vec![0.0; input_data.len()];

        let batch_size = input_shape[0];
        let height = input_shape[2];
        let width = input_shape[3];
        let spatial_size = height * width;

        for b in 0..batch_size {
            for c in 0..channels {
                for spatial in 0..spatial_size {
                    let idx = b * channels * spatial_size + c * spatial_size + spatial;
                    output_data[idx] = (input_data[idx] - self.mean.get(c).unwrap_or(&0.0))
                        / self.std.get(c).unwrap_or(&1.0);
                }
            }
        }

        F32Tensor::from_vec(output_data, input_shape)
    }

    /// 前処理パイプライン（正規化のみ）
    /// Preprocessing pipeline (normalization only)
    pub fn preprocess(&self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // TODO: リサイズ機能の実装
        self.normalize(input)
    }
}

impl Default for F32ImagePreprocessor {
    fn default() -> Self {
        Self::imagenet()
    }
}

// ========================================
// Conv1D/Conv3D - 完全な畳み込み層セット
// Conv1D/Conv3D - Complete convolution layer set
// ========================================

/// 1D畳み込み層
/// 1D Convolution layer
#[derive(Debug)]
pub struct F32Conv1d {
    pub weight: F32Tensor,       // (out_channels, in_channels, kernel_size)
    pub bias: Option<F32Tensor>, // (out_channels,)
    pub weight_grad: Option<F32Tensor>,
    pub bias_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: usize,
    pub stride: usize,
    pub padding: usize,
    pub dilation: usize,
}

impl F32Conv1d {
    /// 新しい1D畳み込み層を作成
    /// Create a new 1D convolution layer
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: usize,
        stride: usize,
        padding: usize,
        dilation: usize,
        bias: bool,
    ) -> Result<Self, RusTorchError> {
        // Heの初期化
        let fan_in = in_channels * kernel_size;
        let std = (2.0 / fan_in as f32).sqrt();

        let weight_shape = vec![out_channels, in_channels, kernel_size];
        let weight = F32Tensor::randn(&weight_shape);
        let std_tensor = F32Tensor::from_scalar(std)?;
        let weight = weight?.mul(&std_tensor)?;

        let bias_tensor = if bias {
            Some(F32Tensor::zeros(&[out_channels])?)
        } else {
            None
        };

        Ok(Self {
            weight,
            bias: bias_tensor,
            weight_grad: None,
            bias_grad: None,
            last_input: None,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            dilation,
        })
    }

    /// フォワードパス（1D畳み込み）
    /// Forward pass (1D convolution)
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 入力を保存（バックワードパス用）
        self.last_input = Some(input.clone());

        // 入力形状: (batch_size, in_channels, length)
        let input_shape = input.shape();
        if input_shape.len() != 3 {
            return Err(RusTorchError::tensor_op(
                "Conv1d input must be 3D (batch, channels, length)",
            ));
        }

        let batch_size = input_shape[0];
        let input_length = input_shape[2];

        // 出力サイズ計算
        let effective_kernel_size = self.dilation * (self.kernel_size - 1) + 1;
        let output_length =
            (input_length + 2 * self.padding - effective_kernel_size) / self.stride + 1;

        // 1D畳み込み実装
        let mut output_data = vec![0.0; batch_size * self.out_channels * output_length];
        let input_data = input.as_slice();
        let weight_data = self.weight.as_slice();

        for b in 0..batch_size {
            for out_c in 0..self.out_channels {
                for ol in 0..output_length {
                    let mut sum = 0.0;

                    // カーネル内の畳み込み
                    for k in 0..self.kernel_size {
                        let il = ol * self.stride + k * self.dilation;

                        // パディングチェック
                        if il >= self.padding && il < input_length + self.padding {
                            let il_actual = il - self.padding;

                            for in_c in 0..self.in_channels {
                                let input_idx = b * self.in_channels * input_length
                                    + in_c * input_length
                                    + il_actual;
                                let weight_idx = out_c * self.in_channels * self.kernel_size
                                    + in_c * self.kernel_size
                                    + k;

                                sum += input_data[input_idx] * weight_data[weight_idx];
                            }
                        }
                    }

                    // バイアス追加
                    if let Some(ref bias) = self.bias {
                        sum += bias.as_slice()[out_c];
                    }

                    let output_idx =
                        b * self.out_channels * output_length + out_c * output_length + ol;
                    output_data[output_idx] = sum;
                }
            }
        }

        let output_shape = vec![batch_size, self.out_channels, output_length];
        F32Tensor::from_vec(output_data, &output_shape)
    }
}

impl F32Layer for F32Conv1d {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            // 重み勾配計算（簡素化）
            self.weight_grad = Some(F32Tensor::zeros(self.weight.shape())?);

            // バイアス勾配計算
            if self.bias.is_some() {
                self.bias_grad = Some(F32Tensor::zeros(&[self.out_channels])?);
            }

            // 入力勾配（簡素化：パススルー）
            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight];
        if let Some(ref bias) = self.bias {
            params.push(bias);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> RusTorchResult<()> {
        // Update weight
        if let Some(ref weight_grad) = self.weight_grad {
            let update = weight_grad.mul_scalar(learning_rate)?;
            self.weight = self.weight.sub(&update)?;
        }

        // Update bias
        if let (Some(ref mut bias), Some(ref bias_grad)) = (&mut self.bias, &self.bias_grad) {
            let update = bias_grad.mul_scalar(learning_rate)?;
            *bias = bias.sub(&update)?;
        }

        Ok(())
    }
}

/// 3D畳み込み層
/// 3D Convolution layer
#[derive(Debug)]
pub struct F32Conv3d {
    pub weight: F32Tensor, // (out_channels, in_channels, kernel_d, kernel_h, kernel_w)
    pub bias: Option<F32Tensor>, // (out_channels,)
    pub weight_grad: Option<F32Tensor>,
    pub bias_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: (usize, usize, usize), // (depth, height, width)
    pub stride: (usize, usize, usize),
    pub padding: (usize, usize, usize),
}

impl F32Conv3d {
    /// 新しい3D畳み込み層を作成
    /// Create a new 3D convolution layer
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: (usize, usize, usize),
        stride: (usize, usize, usize),
        padding: (usize, usize, usize),
        bias: bool,
    ) -> Result<Self, RusTorchError> {
        let (kernel_d, kernel_h, kernel_w) = kernel_size;

        // Heの初期化
        let fan_in = in_channels * kernel_d * kernel_h * kernel_w;
        let std = (2.0 / fan_in as f32).sqrt();

        let weight_shape = vec![out_channels, in_channels, kernel_d, kernel_h, kernel_w];
        let weight = F32Tensor::randn(&weight_shape);
        let std_tensor = F32Tensor::from_scalar(std)?;
        let weight = weight?.mul(&std_tensor)?;

        let bias_tensor = if bias {
            Some(F32Tensor::zeros(&[out_channels])?)
        } else {
            None
        };

        Ok(Self {
            weight,
            bias: bias_tensor,
            weight_grad: None,
            bias_grad: None,
            last_input: None,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
        })
    }

    /// フォワードパス（3D畳み込み）
    /// Forward pass (3D convolution)
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 入力を保存（バックワードパス用）
        self.last_input = Some(input.clone());

        // 入力形状: (batch_size, in_channels, depth, height, width)
        let input_shape = input.shape();
        if input_shape.len() != 5 {
            return Err(RusTorchError::tensor_op(
                "Conv3d input must be 5D (batch, channels, depth, height, width)",
            ));
        }

        let batch_size = input_shape[0];
        let input_depth = input_shape[2];
        let input_height = input_shape[3];
        let input_width = input_shape[4];

        // 出力サイズ計算
        let output_depth =
            (input_depth + 2 * self.padding.0 - self.kernel_size.0) / self.stride.0 + 1;
        let output_height =
            (input_height + 2 * self.padding.1 - self.kernel_size.1) / self.stride.1 + 1;
        let output_width =
            (input_width + 2 * self.padding.2 - self.kernel_size.2) / self.stride.2 + 1;

        // 3D畳み込み実装
        let mut output_data =
            vec![0.0; batch_size * self.out_channels * output_depth * output_height * output_width];
        let input_data = input.as_slice();
        let weight_data = self.weight.as_slice();

        for b in 0..batch_size {
            for out_c in 0..self.out_channels {
                for od in 0..output_depth {
                    for oh in 0..output_height {
                        for ow in 0..output_width {
                            let mut sum = 0.0;

                            // カーネル内の畳み込み
                            for kd in 0..self.kernel_size.0 {
                                for kh in 0..self.kernel_size.1 {
                                    for kw in 0..self.kernel_size.2 {
                                        let id = od * self.stride.0 + kd;
                                        let ih = oh * self.stride.1 + kh;
                                        let iw = ow * self.stride.2 + kw;

                                        // パディングチェック
                                        if id >= self.padding.0
                                            && id < input_depth + self.padding.0
                                            && ih >= self.padding.1
                                            && ih < input_height + self.padding.1
                                            && iw >= self.padding.2
                                            && iw < input_width + self.padding.2
                                        {
                                            let id_actual = id - self.padding.0;
                                            let ih_actual = ih - self.padding.1;
                                            let iw_actual = iw - self.padding.2;

                                            for in_c in 0..self.in_channels {
                                                let input_idx = b
                                                    * self.in_channels
                                                    * input_depth
                                                    * input_height
                                                    * input_width
                                                    + in_c
                                                        * input_depth
                                                        * input_height
                                                        * input_width
                                                    + id_actual * input_height * input_width
                                                    + ih_actual * input_width
                                                    + iw_actual;
                                                let weight_idx = out_c
                                                    * self.in_channels
                                                    * self.kernel_size.0
                                                    * self.kernel_size.1
                                                    * self.kernel_size.2
                                                    + in_c
                                                        * self.kernel_size.0
                                                        * self.kernel_size.1
                                                        * self.kernel_size.2
                                                    + kd * self.kernel_size.1 * self.kernel_size.2
                                                    + kh * self.kernel_size.2
                                                    + kw;

                                                sum +=
                                                    input_data[input_idx] * weight_data[weight_idx];
                                            }
                                        }
                                    }
                                }
                            }

                            // バイアス追加
                            if let Some(ref bias) = self.bias {
                                sum += bias.as_slice()[out_c];
                            }

                            let output_idx =
                                b * self.out_channels * output_depth * output_height * output_width
                                    + out_c * output_depth * output_height * output_width
                                    + od * output_height * output_width
                                    + oh * output_width
                                    + ow;
                            output_data[output_idx] = sum;
                        }
                    }
                }
            }
        }

        let output_shape = vec![
            batch_size,
            self.out_channels,
            output_depth,
            output_height,
            output_width,
        ];
        F32Tensor::from_vec(output_data, &output_shape)
    }
}

impl F32Layer for F32Conv3d {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            // 重み勾配計算（簡素化）
            self.weight_grad = Some(F32Tensor::zeros(self.weight.shape())?);

            // バイアス勾配計算
            if self.bias.is_some() {
                self.bias_grad = Some(F32Tensor::zeros(&[self.out_channels])?);
            }

            // 入力勾配（簡素化：パススルー）
            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight];
        if let Some(ref bias) = self.bias {
            params.push(bias);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> RusTorchResult<()> {
        // Update weight
        if let Some(ref weight_grad) = self.weight_grad {
            let update = weight_grad.mul_scalar(learning_rate)?;
            self.weight = self.weight.sub(&update)?;
        }

        // Update bias
        if let (Some(ref mut bias), Some(ref bias_grad)) = (&mut self.bias, &self.bias_grad) {
            let update = bias_grad.mul_scalar(learning_rate)?;
            *bias = bias.sub(&update)?;
        }

        Ok(())
    }
} // ========================================
  // RNN/LSTM/GRU - 時系列処理層
  // RNN/LSTM/GRU - Time series processing layers
  // ========================================

/// RNN層（Vanilla RNN）
/// RNN layer (Vanilla RNN)
#[derive(Debug)]
pub struct F32RNN {
    pub input_size: usize,
    pub hidden_size: usize,
    pub weight_ih: F32Tensor, // input to hidden weights
    pub weight_hh: F32Tensor, // hidden to hidden weights
    pub bias_ih: Option<F32Tensor>,
    pub bias_hh: Option<F32Tensor>,
    pub weight_ih_grad: Option<F32Tensor>,
    pub weight_hh_grad: Option<F32Tensor>,
    pub bias_ih_grad: Option<F32Tensor>,
    pub bias_hh_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub last_hidden: Option<F32Tensor>,
    pub bidirectional: bool,
    pub batch_first: bool,
}

impl F32RNN {
    /// 新しいRNN層を作成
    /// Create a new RNN layer
    pub fn new(
        input_size: usize,
        hidden_size: usize,
        bias: bool,
        bidirectional: bool,
        batch_first: bool,
    ) -> Result<Self, RusTorchError> {
        // Xavier初期化
        let weight_ih_std = 1.0 / (input_size as f32).sqrt();
        let weight_hh_std = 1.0 / (hidden_size as f32).sqrt();

        let weight_ih = F32Tensor::randn(&[hidden_size, input_size])?
            .mul(&F32Tensor::from_scalar(weight_ih_std)?)?;
        let weight_hh = F32Tensor::randn(&[hidden_size, hidden_size])?
            .mul(&F32Tensor::from_scalar(weight_hh_std)?)?;

        let bias_ih = if bias {
            Some(F32Tensor::zeros(&[hidden_size])?)
        } else {
            None
        };

        let bias_hh = if bias {
            Some(F32Tensor::zeros(&[hidden_size])?)
        } else {
            None
        };

        Ok(Self {
            input_size,
            hidden_size,
            weight_ih,
            weight_hh,
            bias_ih,
            bias_hh,
            weight_ih_grad: None,
            weight_hh_grad: None,
            bias_ih_grad: None,
            bias_hh_grad: None,
            last_input: None,
            last_hidden: None,
            bidirectional,
            batch_first,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        input: &F32Tensor,
        hidden: Option<&F32Tensor>,
    ) -> Result<(F32Tensor, F32Tensor), RusTorchError> {
        self.last_input = Some(input.clone());

        let input_shape = input.shape();
        let (batch_size, seq_len) = if self.batch_first {
            (input_shape[0], input_shape[1])
        } else {
            (input_shape[1], input_shape[0])
        };

        // 初期隠れ状態
        let mut h = if let Some(h) = hidden {
            h.clone()
        } else {
            F32Tensor::zeros(&[batch_size, self.hidden_size])?
        };

        let mut outputs = Vec::new();

        // 各時刻でのRNN計算
        for t in 0..seq_len {
            let x_t = if self.batch_first {
                input
                    .slice(&[(0, batch_size), (t, t + 1), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            } else {
                input
                    .slice(&[(t, t + 1), (0, batch_size), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            };

            // RNN cell: h_t = tanh(W_ih @ x_t + W_hh @ h_{t-1} + b)
            let ih = x_t.matmul(&self.weight_ih.transpose()?)?;
            let hh = h.matmul(&self.weight_hh.transpose()?)?;
            let mut gate = ih.add(&hh)?;

            if let Some(ref bias_ih) = self.bias_ih {
                let bias_ih_expanded = bias_ih
                    .unsqueeze(0)?
                    .expand(&[batch_size, self.hidden_size])?;
                gate = gate.add(&bias_ih_expanded)?;
            }

            if let Some(ref bias_hh) = self.bias_hh {
                let bias_hh_expanded = bias_hh
                    .unsqueeze(0)?
                    .expand(&[batch_size, self.hidden_size])?;
                gate = gate.add(&bias_hh_expanded)?;
            }

            h = F32Activation::Tanh.forward(&gate)?;
            outputs.push(h.clone());
        }

        // 出力を結合
        let output = if self.batch_first {
            // (batch, seq_len, hidden_size)
            let mut output_data = Vec::new();
            for output_t in &outputs {
                output_data.extend_from_slice(output_t.as_slice());
            }
            F32Tensor::from_vec(output_data, &[batch_size, seq_len, self.hidden_size])?
        } else {
            // (seq_len, batch, hidden_size)
            let mut output_data = Vec::new();
            for t in 0..seq_len {
                output_data.extend_from_slice(outputs[t].as_slice());
            }
            F32Tensor::from_vec(output_data, &[seq_len, batch_size, self.hidden_size])?
        };

        self.last_hidden = Some(h.clone());
        Ok((output, h))
    }
}

impl F32Layer for F32RNN {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let (output, _) = self.forward(input, None)?;
        Ok(output)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            // 勾配を初期化
            self.weight_ih_grad = Some(F32Tensor::zeros(self.weight_ih.shape())?);
            self.weight_hh_grad = Some(F32Tensor::zeros(self.weight_hh.shape())?);

            if self.bias_ih.is_some() {
                self.bias_ih_grad = Some(F32Tensor::zeros(&[self.hidden_size])?);
            }
            if self.bias_hh.is_some() {
                self.bias_hh_grad = Some(F32Tensor::zeros(&[self.hidden_size])?);
            }

            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight_ih, &self.weight_hh];
        if let Some(ref bias_ih) = self.bias_ih {
            params.push(bias_ih);
        }
        if let Some(ref bias_hh) = self.bias_hh {
            params.push(bias_hh);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        let lr_tensor = F32Tensor::from_scalar(learning_rate)?;

        if let Some(ref weight_ih_grad) = self.weight_ih_grad {
            let update = weight_ih_grad.mul(&lr_tensor)?;
            self.weight_ih = self.weight_ih.sub(&update)?;
        }

        if let Some(ref weight_hh_grad) = self.weight_hh_grad {
            let update = weight_hh_grad.mul(&lr_tensor)?;
            self.weight_hh = self.weight_hh.sub(&update)?;
        }

        if let (Some(ref mut bias_ih), Some(ref bias_ih_grad)) =
            (&mut self.bias_ih, &self.bias_ih_grad)
        {
            let update = bias_ih_grad.mul(&lr_tensor)?;
            *bias_ih = bias_ih.sub(&update)?;
        }

        if let (Some(ref mut bias_hh), Some(ref bias_hh_grad)) =
            (&mut self.bias_hh, &self.bias_hh_grad)
        {
            let update = bias_hh_grad.mul(&lr_tensor)?;
            *bias_hh = bias_hh.sub(&update)?;
        }

        Ok(())
    }
}

/// LSTM層（Long Short-Term Memory）
/// LSTM layer (Long Short-Term Memory)
#[derive(Debug)]
pub struct F32LSTM {
    pub input_size: usize,
    pub hidden_size: usize,
    // 4つのゲート用の重み: input, forget, cell, output
    pub weight_ih: F32Tensor,       // (4 * hidden_size, input_size)
    pub weight_hh: F32Tensor,       // (4 * hidden_size, hidden_size)
    pub bias_ih: Option<F32Tensor>, // (4 * hidden_size)
    pub bias_hh: Option<F32Tensor>, // (4 * hidden_size)
    pub weight_ih_grad: Option<F32Tensor>,
    pub weight_hh_grad: Option<F32Tensor>,
    pub bias_ih_grad: Option<F32Tensor>,
    pub bias_hh_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub last_hidden: Option<F32Tensor>,
    pub last_cell: Option<F32Tensor>,
    pub bidirectional: bool,
    pub batch_first: bool,
}

impl F32LSTM {
    /// 新しいLSTM層を作成
    /// Create a new LSTM layer
    pub fn new(
        input_size: usize,
        hidden_size: usize,
        bias: bool,
        bidirectional: bool,
        batch_first: bool,
    ) -> Result<Self, RusTorchError> {
        // Xavier初期化
        let weight_ih_std = 1.0 / (input_size as f32).sqrt();
        let weight_hh_std = 1.0 / (hidden_size as f32).sqrt();

        let weight_ih = F32Tensor::randn(&[4 * hidden_size, input_size])?
            .mul(&F32Tensor::from_scalar(weight_ih_std)?)?;
        let weight_hh = F32Tensor::randn(&[4 * hidden_size, hidden_size])?
            .mul(&F32Tensor::from_scalar(weight_hh_std)?)?;

        let bias_ih = if bias {
            // forget gate biasを1.0で初期化（一般的な慣習）
            let mut bias_data = vec![0.0; 4 * hidden_size];
            for i in hidden_size..(2 * hidden_size) {
                bias_data[i] = 1.0; // forget gate bias
            }
            Some(F32Tensor::from_vec(bias_data, &[4 * hidden_size])?)
        } else {
            None
        };

        let bias_hh = if bias {
            Some(F32Tensor::zeros(&[4 * hidden_size])?)
        } else {
            None
        };

        Ok(Self {
            input_size,
            hidden_size,
            weight_ih,
            weight_hh,
            bias_ih,
            bias_hh,
            weight_ih_grad: None,
            weight_hh_grad: None,
            bias_ih_grad: None,
            bias_hh_grad: None,
            last_input: None,
            last_hidden: None,
            last_cell: None,
            bidirectional,
            batch_first,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        input: &F32Tensor,
        state: Option<(&F32Tensor, &F32Tensor)>,
    ) -> Result<(F32Tensor, F32Tensor, F32Tensor), RusTorchError> {
        self.last_input = Some(input.clone());

        let input_shape = input.shape();
        let (batch_size, seq_len) = if self.batch_first {
            (input_shape[0], input_shape[1])
        } else {
            (input_shape[1], input_shape[0])
        };

        // 初期状態
        let (mut h, mut c) = if let Some((h, c)) = state {
            (h.clone(), c.clone())
        } else {
            (
                F32Tensor::zeros(&[batch_size, self.hidden_size])?,
                F32Tensor::zeros(&[batch_size, self.hidden_size])?,
            )
        };

        let mut outputs = Vec::new();

        // 各時刻でのLSTM計算
        for t in 0..seq_len {
            let x_t = if self.batch_first {
                input
                    .slice(&[(0, batch_size), (t, t + 1), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            } else {
                input
                    .slice(&[(t, t + 1), (0, batch_size), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            };

            // LSTM cell computation
            let ih = x_t.matmul(&self.weight_ih.transpose()?)?;
            let hh = h.matmul(&self.weight_hh.transpose()?)?;
            let mut gates = ih.add(&hh)?;

            if let Some(ref bias_ih) = self.bias_ih {
                let bias_ih_expanded = bias_ih
                    .unsqueeze(0)?
                    .expand(&[batch_size, 4 * self.hidden_size])?;
                gates = gates.add(&bias_ih_expanded)?;
            }

            if let Some(ref bias_hh) = self.bias_hh {
                let bias_hh_expanded = bias_hh
                    .unsqueeze(0)?
                    .expand(&[batch_size, 4 * self.hidden_size])?;
                gates = gates.add(&bias_hh_expanded)?;
            }

            // 4つのゲートに分割
            let i_gate = gates.slice(&[(0, batch_size), (0, self.hidden_size)])?;
            let f_gate =
                gates.slice(&[(0, batch_size), (self.hidden_size, 2 * self.hidden_size)])?;
            let g_gate = gates.slice(&[
                (0, batch_size),
                (2 * self.hidden_size, 3 * self.hidden_size),
            ])?;
            let o_gate = gates.slice(&[
                (0, batch_size),
                (3 * self.hidden_size, 4 * self.hidden_size),
            ])?;

            // ゲート活性化
            let i = F32Activation::Sigmoid.forward(&i_gate)?;
            let f = F32Activation::Sigmoid.forward(&f_gate)?;
            let g = F32Activation::Tanh.forward(&g_gate)?;
            let o = F32Activation::Sigmoid.forward(&o_gate)?;

            // セル状態更新
            c = f.mul(&c)?.add(&i.mul(&g)?)?;

            // 隠れ状態更新
            let c_tanh = F32Activation::Tanh.forward(&c)?;
            h = o.mul(&c_tanh)?;

            outputs.push(h.clone());
        }

        // 出力を結合
        let output = if self.batch_first {
            let mut output_data = Vec::new();
            for output_t in &outputs {
                output_data.extend_from_slice(output_t.as_slice());
            }
            F32Tensor::from_vec(output_data, &[batch_size, seq_len, self.hidden_size])?
        } else {
            let mut output_data = Vec::new();
            for t in 0..seq_len {
                output_data.extend_from_slice(outputs[t].as_slice());
            }
            F32Tensor::from_vec(output_data, &[seq_len, batch_size, self.hidden_size])?
        };

        self.last_hidden = Some(h.clone());
        self.last_cell = Some(c.clone());
        Ok((output, h, c))
    }
}

impl F32Layer for F32LSTM {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let (output, _, _) = self.forward(input, None)?;
        Ok(output)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            self.weight_ih_grad = Some(F32Tensor::zeros(self.weight_ih.shape())?);
            self.weight_hh_grad = Some(F32Tensor::zeros(self.weight_hh.shape())?);

            if self.bias_ih.is_some() {
                self.bias_ih_grad = Some(F32Tensor::zeros(&[4 * self.hidden_size])?);
            }
            if self.bias_hh.is_some() {
                self.bias_hh_grad = Some(F32Tensor::zeros(&[4 * self.hidden_size])?);
            }

            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight_ih, &self.weight_hh];
        if let Some(ref bias_ih) = self.bias_ih {
            params.push(bias_ih);
        }
        if let Some(ref bias_hh) = self.bias_hh {
            params.push(bias_hh);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        let lr_tensor = F32Tensor::from_scalar(learning_rate)?;

        if let Some(ref weight_ih_grad) = self.weight_ih_grad {
            let update = weight_ih_grad.mul(&lr_tensor)?;
            self.weight_ih = self.weight_ih.sub(&update)?;
        }

        if let Some(ref weight_hh_grad) = self.weight_hh_grad {
            let update = weight_hh_grad.mul(&lr_tensor)?;
            self.weight_hh = self.weight_hh.sub(&update)?;
        }

        if let (Some(ref mut bias_ih), Some(ref bias_ih_grad)) =
            (&mut self.bias_ih, &self.bias_ih_grad)
        {
            let update = bias_ih_grad.mul(&lr_tensor)?;
            *bias_ih = bias_ih.sub(&update)?;
        }

        if let (Some(ref mut bias_hh), Some(ref bias_hh_grad)) =
            (&mut self.bias_hh, &self.bias_hh_grad)
        {
            let update = bias_hh_grad.mul(&lr_tensor)?;
            *bias_hh = bias_hh.sub(&update)?;
        }

        Ok(())
    }
}

/// GRU層（Gated Recurrent Unit）
/// GRU layer (Gated Recurrent Unit)
#[derive(Debug)]
pub struct F32GRU {
    pub input_size: usize,
    pub hidden_size: usize,
    // 3つのゲート用の重み: reset, update, new
    pub weight_ih: F32Tensor,       // (3 * hidden_size, input_size)
    pub weight_hh: F32Tensor,       // (3 * hidden_size, hidden_size)
    pub bias_ih: Option<F32Tensor>, // (3 * hidden_size)
    pub bias_hh: Option<F32Tensor>, // (3 * hidden_size)
    pub weight_ih_grad: Option<F32Tensor>,
    pub weight_hh_grad: Option<F32Tensor>,
    pub bias_ih_grad: Option<F32Tensor>,
    pub bias_hh_grad: Option<F32Tensor>,
    pub last_input: Option<F32Tensor>,
    pub last_hidden: Option<F32Tensor>,
    pub bidirectional: bool,
    pub batch_first: bool,
}

impl F32GRU {
    /// 新しいGRU層を作成
    /// Create a new GRU layer
    pub fn new(
        input_size: usize,
        hidden_size: usize,
        bias: bool,
        bidirectional: bool,
        batch_first: bool,
    ) -> Result<Self, RusTorchError> {
        // Xavier初期化
        let weight_ih_std = 1.0 / (input_size as f32).sqrt();
        let weight_hh_std = 1.0 / (hidden_size as f32).sqrt();

        let weight_ih = F32Tensor::randn(&[3 * hidden_size, input_size])?
            .mul(&F32Tensor::from_scalar(weight_ih_std)?)?;
        let weight_hh = F32Tensor::randn(&[3 * hidden_size, hidden_size])?
            .mul(&F32Tensor::from_scalar(weight_hh_std)?)?;

        let bias_ih = if bias {
            Some(F32Tensor::zeros(&[3 * hidden_size])?)
        } else {
            None
        };

        let bias_hh = if bias {
            Some(F32Tensor::zeros(&[3 * hidden_size])?)
        } else {
            None
        };

        Ok(Self {
            input_size,
            hidden_size,
            weight_ih,
            weight_hh,
            bias_ih,
            bias_hh,
            weight_ih_grad: None,
            weight_hh_grad: None,
            bias_ih_grad: None,
            bias_hh_grad: None,
            last_input: None,
            last_hidden: None,
            bidirectional,
            batch_first,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        input: &F32Tensor,
        hidden: Option<&F32Tensor>,
    ) -> Result<(F32Tensor, F32Tensor), RusTorchError> {
        self.last_input = Some(input.clone());

        let input_shape = input.shape();
        let (batch_size, seq_len) = if self.batch_first {
            (input_shape[0], input_shape[1])
        } else {
            (input_shape[1], input_shape[0])
        };

        // 初期隠れ状態
        let mut h = if let Some(h) = hidden {
            h.clone()
        } else {
            F32Tensor::zeros(&[batch_size, self.hidden_size])?
        };

        let mut outputs = Vec::new();

        // 各時刻でのGRU計算
        for t in 0..seq_len {
            let x_t = if self.batch_first {
                input
                    .slice(&[(0, batch_size), (t, t + 1), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            } else {
                input
                    .slice(&[(t, t + 1), (0, batch_size), (0, self.input_size)])?
                    .reshape(&[batch_size, self.input_size])?
            };

            // GRU cell computation
            let ih = x_t.matmul(&self.weight_ih.transpose()?)?;
            let hh = h.matmul(&self.weight_hh.transpose()?)?;

            let mut gi = ih.clone();
            let mut gh = hh.clone();

            if let Some(ref bias_ih) = self.bias_ih {
                let bias_ih_expanded = bias_ih
                    .unsqueeze(0)?
                    .expand(&[batch_size, 3 * self.hidden_size])?;
                gi = gi.add(&bias_ih_expanded)?;
            }

            if let Some(ref bias_hh) = self.bias_hh {
                let bias_hh_expanded = bias_hh
                    .unsqueeze(0)?
                    .expand(&[batch_size, 3 * self.hidden_size])?;
                gh = gh.add(&bias_hh_expanded)?;
            }

            // reset gate, update gate用の計算
            let i_r = gi.slice(&[(0, batch_size), (0, self.hidden_size)])?;
            let i_z = gi.slice(&[(0, batch_size), (self.hidden_size, 2 * self.hidden_size)])?;
            let h_r = gh.slice(&[(0, batch_size), (0, self.hidden_size)])?;
            let h_z = gh.slice(&[(0, batch_size), (self.hidden_size, 2 * self.hidden_size)])?;

            let r = F32Activation::Sigmoid.forward(&i_r.add(&h_r)?)?;
            let z = F32Activation::Sigmoid.forward(&i_z.add(&h_z)?)?;

            // new gate計算
            let i_n = gi.slice(&[
                (0, batch_size),
                (2 * self.hidden_size, 3 * self.hidden_size),
            ])?;
            let h_n = gh.slice(&[
                (0, batch_size),
                (2 * self.hidden_size, 3 * self.hidden_size),
            ])?;
            let n = F32Activation::Tanh.forward(&i_n.add(&r.mul(&h_n)?)?)?;

            // 隠れ状態更新
            let one = F32Tensor::ones(&[batch_size, self.hidden_size])?;
            let one_minus_z = one.sub(&z)?;
            h = one_minus_z.mul(&n)?.add(&z.mul(&h)?)?;

            outputs.push(h.clone());
        }

        // 出力を結合
        let output = if self.batch_first {
            let mut output_data = Vec::new();
            for output_t in &outputs {
                output_data.extend_from_slice(output_t.as_slice());
            }
            F32Tensor::from_vec(output_data, &[batch_size, seq_len, self.hidden_size])?
        } else {
            let mut output_data = Vec::new();
            for t in 0..seq_len {
                output_data.extend_from_slice(outputs[t].as_slice());
            }
            F32Tensor::from_vec(output_data, &[seq_len, batch_size, self.hidden_size])?
        };

        self.last_hidden = Some(h.clone());
        Ok((output, h))
    }
}

impl F32Layer for F32GRU {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let (output, _) = self.forward(input, None)?;
        Ok(output)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            self.weight_ih_grad = Some(F32Tensor::zeros(self.weight_ih.shape())?);
            self.weight_hh_grad = Some(F32Tensor::zeros(self.weight_hh.shape())?);

            if self.bias_ih.is_some() {
                self.bias_ih_grad = Some(F32Tensor::zeros(&[3 * self.hidden_size])?);
            }
            if self.bias_hh.is_some() {
                self.bias_hh_grad = Some(F32Tensor::zeros(&[3 * self.hidden_size])?);
            }

            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = vec![&self.weight_ih, &self.weight_hh];
        if let Some(ref bias_ih) = self.bias_ih {
            params.push(bias_ih);
        }
        if let Some(ref bias_hh) = self.bias_hh {
            params.push(bias_hh);
        }
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        let lr_tensor = F32Tensor::from_scalar(learning_rate)?;

        if let Some(ref weight_ih_grad) = self.weight_ih_grad {
            let update = weight_ih_grad.mul(&lr_tensor)?;
            self.weight_ih = self.weight_ih.sub(&update)?;
        }

        if let Some(ref weight_hh_grad) = self.weight_hh_grad {
            let update = weight_hh_grad.mul(&lr_tensor)?;
            self.weight_hh = self.weight_hh.sub(&update)?;
        }

        if let (Some(ref mut bias_ih), Some(ref bias_ih_grad)) =
            (&mut self.bias_ih, &self.bias_ih_grad)
        {
            let update = bias_ih_grad.mul(&lr_tensor)?;
            *bias_ih = bias_ih.sub(&update)?;
        }

        if let (Some(ref mut bias_hh), Some(ref bias_hh_grad)) =
            (&mut self.bias_hh, &self.bias_hh_grad)
        {
            let update = bias_hh_grad.mul(&lr_tensor)?;
            *bias_hh = bias_hh.sub(&update)?;
        }

        Ok(())
    }
} // ========================================
  // Transformer - 現代的アーキテクチャ
  // Transformer - Modern architecture
  // ========================================

/// 位置エンコーディング
/// Positional Encoding
#[derive(Debug)]
pub struct F32PositionalEncoding {
    pub encoding: F32Tensor,
    pub max_len: usize,
    pub d_model: usize,
}

impl F32PositionalEncoding {
    /// 新しい位置エンコーディングを作成
    /// Create new positional encoding
    pub fn new(d_model: usize, max_len: usize) -> Result<Self, RusTorchError> {
        let mut pe_data = vec![0.0; max_len * d_model];

        for pos in 0..max_len {
            for i in (0..d_model).step_by(2) {
                let angle = pos as f32 / 10000.0_f32.powf(i as f32 / d_model as f32);
                pe_data[pos * d_model + i] = angle.sin();
                if i + 1 < d_model {
                    pe_data[pos * d_model + i + 1] = angle.cos();
                }
            }
        }

        let encoding = F32Tensor::from_vec(pe_data, &[max_len, d_model])?;

        Ok(Self {
            encoding,
            max_len,
            d_model,
        })
    }

    /// 位置エンコーディングを適用
    /// Apply positional encoding
    pub fn forward(&self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let input_shape = input.shape();
        let seq_len = input_shape[1];

        if seq_len > self.max_len {
            return Err(RusTorchError::tensor_op("Sequence length exceeds max_len"));
        }

        let pe_slice = self.encoding.slice(&[(0, seq_len), (0, self.d_model)])?;
        let pe_expanded =
            pe_slice
                .unsqueeze(0)?
                .expand(&[input_shape[0], seq_len, self.d_model])?;

        input.add(&pe_expanded)
    }
}

/// マルチヘッドアテンション
/// Multi-Head Attention
#[derive(Debug)]
pub struct F32MultiHeadAttention {
    pub d_model: usize,
    pub num_heads: usize,
    pub d_k: usize,
    pub d_v: usize,
    pub w_q: F32Tensor,
    pub w_k: F32Tensor,
    pub w_v: F32Tensor,
    pub w_o: F32Tensor,
    pub dropout_rate: f32,
    pub last_input: Option<F32Tensor>,
    pub weight_grads: Option<(F32Tensor, F32Tensor, F32Tensor, F32Tensor)>,
}

impl F32MultiHeadAttention {
    /// 新しいマルチヘッドアテンションを作成
    /// Create new multi-head attention
    pub fn new(d_model: usize, num_heads: usize, dropout_rate: f32) -> Result<Self, RusTorchError> {
        assert_eq!(
            d_model % num_heads,
            0,
            "d_model must be divisible by num_heads"
        );

        let d_k = d_model / num_heads;
        let d_v = d_model / num_heads;

        // Xavier初期化
        let std = 1.0 / (d_model as f32).sqrt();

        let w_q = F32Tensor::randn(&[d_model, d_model])?.mul(&F32Tensor::from_scalar(std)?)?;
        let w_k = F32Tensor::randn(&[d_model, d_model])?.mul(&F32Tensor::from_scalar(std)?)?;
        let w_v = F32Tensor::randn(&[d_model, d_model])?.mul(&F32Tensor::from_scalar(std)?)?;
        let w_o = F32Tensor::randn(&[d_model, d_model])?.mul(&F32Tensor::from_scalar(std)?)?;

        Ok(Self {
            d_model,
            num_heads,
            d_k,
            d_v,
            w_q,
            w_k,
            w_v,
            w_o,
            dropout_rate,
            last_input: None,
            weight_grads: None,
        })
    }

    /// スケールドドットプロダクトアテンション
    /// Scaled Dot-Product Attention
    fn scaled_dot_product_attention(
        &self,
        q: &F32Tensor,
        k: &F32Tensor,
        v: &F32Tensor,
        mask: Option<&F32Tensor>,
    ) -> Result<F32Tensor, RusTorchError> {
        let d_k = self.d_k as f32;
        let scale = 1.0 / d_k.sqrt();

        // Attention scores: Q @ K^T / sqrt(d_k)
        let scores = q
            .matmul(&k.transpose()?)?
            .mul(&F32Tensor::from_scalar(scale)?)?;

        // マスクを適用（オプション）
        let scores = if let Some(mask) = mask {
            scores.add(&mask)?
        } else {
            scores
        };

        // Softmax
        let attention_weights = F32Activation::Softmax.forward(&scores)?;

        // アテンション重みをVに適用
        attention_weights.matmul(v)
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        query: &F32Tensor,
        key: &F32Tensor,
        value: &F32Tensor,
        mask: Option<&F32Tensor>,
    ) -> Result<F32Tensor, RusTorchError> {
        self.last_input = Some(query.clone());

        let batch_size = query.shape()[0];
        let seq_len = query.shape()[1];

        // Q, K, Vを計算
        let q = query.matmul(&self.w_q)?;
        let k = key.matmul(&self.w_k)?;
        let v = value.matmul(&self.w_v)?;

        // ヘッドに分割: (batch_size, seq_len, d_model) -> (batch_size, num_heads, seq_len, d_k)
        let q_heads = q
            .reshape(&[batch_size, seq_len, self.num_heads, self.d_k])?
            .transpose_dims(1, 2)?;
        let k_heads = k
            .reshape(&[batch_size, seq_len, self.num_heads, self.d_k])?
            .transpose_dims(1, 2)?;
        let v_heads = v
            .reshape(&[batch_size, seq_len, self.num_heads, self.d_v])?
            .transpose_dims(1, 2)?;

        // 各ヘッドでアテンション計算
        let mut head_outputs = Vec::new();
        for h in 0..self.num_heads {
            let q_h = q_heads
                .slice(&[(0, batch_size), (h, h + 1), (0, seq_len), (0, self.d_k)])?
                .reshape(&[batch_size, seq_len, self.d_k])?;
            let k_h = k_heads
                .slice(&[(0, batch_size), (h, h + 1), (0, seq_len), (0, self.d_k)])?
                .reshape(&[batch_size, seq_len, self.d_k])?;
            let v_h = v_heads
                .slice(&[(0, batch_size), (h, h + 1), (0, seq_len), (0, self.d_v)])?
                .reshape(&[batch_size, seq_len, self.d_v])?;

            let head_output = self.scaled_dot_product_attention(&q_h, &k_h, &v_h, mask)?;
            head_outputs.push(head_output);
        }

        // ヘッドを結合
        let mut concat_data = Vec::new();
        for i in 0..batch_size {
            for j in 0..seq_len {
                for head in &head_outputs {
                    let head_data = head.slice(&[(i, i + 1), (j, j + 1), (0, self.d_v)])?;
                    concat_data.extend_from_slice(head_data.as_slice());
                }
            }
        }

        let concatenated = F32Tensor::from_vec(concat_data, &[batch_size, seq_len, self.d_model])?;

        // 出力射影
        concatenated.matmul(&self.w_o)
    }
}

impl F32Layer for F32MultiHeadAttention {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // self-attentionの場合
        self.forward(input, input, input, None)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        if let Some(ref last_input) = self.last_input {
            // 重み勾配を初期化
            let w_q_grad = F32Tensor::zeros(self.w_q.shape())?;
            let w_k_grad = F32Tensor::zeros(self.w_k.shape())?;
            let w_v_grad = F32Tensor::zeros(self.w_v.shape())?;
            let w_o_grad = F32Tensor::zeros(self.w_o.shape())?;

            self.weight_grads = Some((w_q_grad, w_k_grad, w_v_grad, w_o_grad));

            Ok(last_input.clone())
        } else {
            Err(RusTorchError::tensor_op("No saved input for backward pass"))
        }
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        vec![&self.w_q, &self.w_k, &self.w_v, &self.w_o]
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        if let Some(ref grads) = self.weight_grads {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;

            let update_q = grads.0.mul(&lr_tensor)?;
            self.w_q = self.w_q.sub(&update_q)?;

            let update_k = grads.1.mul(&lr_tensor)?;
            self.w_k = self.w_k.sub(&update_k)?;

            let update_v = grads.2.mul(&lr_tensor)?;
            self.w_v = self.w_v.sub(&update_v)?;

            let update_o = grads.3.mul(&lr_tensor)?;
            self.w_o = self.w_o.sub(&update_o)?;
        }
        Ok(())
    }
}

/// フィードフォワードネットワーク
/// Feed Forward Network
#[derive(Debug)]
pub struct F32FeedForward {
    pub linear1: F32Linear,
    pub linear2: F32Linear,
    pub dropout_rate: f32,
}

impl F32FeedForward {
    /// 新しいフィードフォワードネットワークを作成
    /// Create new feed forward network
    pub fn new(d_model: usize, d_ff: usize, dropout_rate: f32) -> Result<Self, RusTorchError> {
        let linear1 = F32Linear::new(d_model, d_ff, true)?;
        let linear2 = F32Linear::new(d_ff, d_model, true)?;

        Ok(Self {
            linear1,
            linear2,
            dropout_rate,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let x = self.linear1.forward(input)?;
        let x = F32Activation::ReLU.forward(&x)?;
        // ドロップアウトは簡素化のため省略
        self.linear2.forward(&x)
    }
}

impl F32Layer for F32FeedForward {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        let grad1 = self.linear2.backward(grad_output)?;
        let grad2 = self.linear1.backward(&grad1)?;
        Ok(grad2)
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = Vec::new();
        params.extend(self.linear1.parameters());
        params.extend(self.linear2.parameters());
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        self.linear1.update_parameters(learning_rate)?;
        self.linear2.update_parameters(learning_rate)?;
        Ok(())
    }
}

/// レイヤー正規化
/// Layer Normalization
#[derive(Debug)]
pub struct F32LayerNorm {
    pub normalized_shape: Vec<usize>,
    pub weight: F32Tensor,
    pub bias: F32Tensor,
    pub eps: f32,
    pub weight_grad: Option<F32Tensor>,
    pub bias_grad: Option<F32Tensor>,
}

impl F32LayerNorm {
    /// 新しいレイヤー正規化を作成
    /// Create new layer normalization
    pub fn new(normalized_shape: Vec<usize>, eps: f32) -> Result<Self, RusTorchError> {
        let num_features = normalized_shape.iter().product();
        let weight = F32Tensor::ones(&[num_features])?;
        let bias = F32Tensor::zeros(&[num_features])?;

        Ok(Self {
            normalized_shape,
            weight,
            bias,
            eps,
            weight_grad: None,
            bias_grad: None,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        let input_shape = input.shape();
        let last_dim = input_shape[input_shape.len() - 1];

        // 最後の次元で正規化
        let input_data = input.as_slice();
        let mut output_data = vec![0.0; input_data.len()];

        let batch_size = input_data.len() / last_dim;

        for batch in 0..batch_size {
            let start_idx = batch * last_dim;
            let end_idx = start_idx + last_dim;
            let batch_data = &input_data[start_idx..end_idx];

            // 平均と分散を計算
            let mean = batch_data.iter().sum::<f32>() / last_dim as f32;
            let variance = batch_data
                .iter()
                .map(|x| (x - mean) * (x - mean))
                .sum::<f32>()
                / last_dim as f32;

            let std = (variance + self.eps).sqrt();

            // 正規化
            let weight_data = self.weight.as_slice();
            let bias_data = self.bias.as_slice();

            for i in 0..last_dim {
                let normalized = (batch_data[i] - mean) / std;
                output_data[start_idx + i] = weight_data[i] * normalized + bias_data[i];
            }
        }

        F32Tensor::from_vec(output_data, input_shape)
    }
}

impl F32Layer for F32LayerNorm {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        self.weight_grad = Some(F32Tensor::zeros(self.weight.shape())?);
        self.bias_grad = Some(F32Tensor::zeros(self.bias.shape())?);
        Ok(grad_output.clone())
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        vec![&self.weight, &self.bias]
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        if let Some(ref weight_grad) = self.weight_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = weight_grad.mul(&lr_tensor)?;
            self.weight = self.weight.sub(&update)?;
        }

        if let Some(ref bias_grad) = self.bias_grad {
            let lr_tensor = F32Tensor::from_scalar(learning_rate)?;
            let update = bias_grad.mul(&lr_tensor)?;
            self.bias = self.bias.sub(&update)?;
        }

        Ok(())
    }
}

/// トランスフォーマーブロック
/// Transformer Block
#[derive(Debug)]
pub struct F32TransformerBlock {
    pub attention: F32MultiHeadAttention,
    pub feed_forward: F32FeedForward,
    pub norm1: F32LayerNorm,
    pub norm2: F32LayerNorm,
    pub dropout_rate: f32,
}

impl F32TransformerBlock {
    /// 新しいトランスフォーマーブロックを作成
    /// Create new transformer block
    pub fn new(
        d_model: usize,
        num_heads: usize,
        d_ff: usize,
        dropout_rate: f32,
    ) -> Result<Self, RusTorchError> {
        let attention = F32MultiHeadAttention::new(d_model, num_heads, dropout_rate)?;
        let feed_forward = F32FeedForward::new(d_model, d_ff, dropout_rate)?;
        let norm1 = F32LayerNorm::new(vec![d_model], 1e-6)?;
        let norm2 = F32LayerNorm::new(vec![d_model], 1e-6)?;

        Ok(Self {
            attention,
            feed_forward,
            norm1,
            norm2,
            dropout_rate,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        input: &F32Tensor,
        mask: Option<&F32Tensor>,
    ) -> Result<F32Tensor, RusTorchError> {
        // Self-attention with residual connection
        let attn_output = self.attention.forward(input, input, input, mask)?;
        let x = input.add(&attn_output)?; // Residual connection
        let x = self.norm1.forward(&x)?;

        // Feed forward with residual connection
        let ff_output = self.feed_forward.forward(&x)?;
        let x = x.add(&ff_output)?; // Residual connection
        let x = self.norm2.forward(&x)?;

        Ok(x)
    }
}

impl F32Layer for F32TransformerBlock {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input, None)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        let grad2 = self.norm2.backward(grad_output)?;
        let grad_ff = self.feed_forward.backward(&grad2)?;
        let grad1 = self.norm1.backward(&grad_ff)?;
        let grad_attn = self.attention.backward(&grad1)?;
        Ok(grad_attn)
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = Vec::new();
        params.extend(self.attention.parameters());
        params.extend(self.feed_forward.parameters());
        params.extend(self.norm1.parameters());
        params.extend(self.norm2.parameters());
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        self.attention.update_parameters(learning_rate)?;
        self.feed_forward.update_parameters(learning_rate)?;
        self.norm1.update_parameters(learning_rate)?;
        self.norm2.update_parameters(learning_rate)?;
        Ok(())
    }
}

/// 完全なトランスフォーマーモデル
/// Complete Transformer Model
#[derive(Debug)]
pub struct F32Transformer {
    pub d_model: usize,
    pub num_heads: usize,
    pub num_layers: usize,
    pub d_ff: usize,
    pub max_seq_len: usize,
    pub vocab_size: usize,
    pub embedding: F32Linear,
    pub positional_encoding: F32PositionalEncoding,
    pub transformer_blocks: Vec<F32TransformerBlock>,
    pub output_projection: F32Linear,
    pub dropout_rate: f32,
}

impl F32Transformer {
    /// 新しいトランスフォーマーモデルを作成
    /// Create new transformer model
    pub fn new(
        vocab_size: usize,
        d_model: usize,
        num_heads: usize,
        num_layers: usize,
        d_ff: usize,
        max_seq_len: usize,
        dropout_rate: f32,
    ) -> Result<Self, RusTorchError> {
        let embedding = F32Linear::new(vocab_size, d_model, false)?;
        let positional_encoding = F32PositionalEncoding::new(d_model, max_seq_len)?;

        let mut transformer_blocks = Vec::new();
        for _ in 0..num_layers {
            let block = F32TransformerBlock::new(d_model, num_heads, d_ff, dropout_rate)?;
            transformer_blocks.push(block);
        }

        let output_projection = F32Linear::new(d_model, vocab_size, true)?;

        Ok(Self {
            d_model,
            num_heads,
            num_layers,
            d_ff,
            max_seq_len,
            vocab_size,
            embedding,
            positional_encoding,
            transformer_blocks,
            output_projection,
            dropout_rate,
        })
    }

    /// フォワードパス
    /// Forward pass
    pub fn forward(
        &mut self,
        input_ids: &F32Tensor,
        mask: Option<&F32Tensor>,
    ) -> Result<F32Tensor, RusTorchError> {
        // Embedding
        let x = self.embedding.forward(input_ids)?;

        // 位置エンコーディングを追加
        let x = self.positional_encoding.forward(&x)?;

        // トランスフォーマーブロックを通す
        let mut x = x;
        for block in &mut self.transformer_blocks {
            x = block.forward(&x, mask)?;
        }

        // 出力投影
        self.output_projection.forward(&x)
    }
}

impl F32Layer for F32Transformer {
    fn forward(&mut self, input: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        self.forward(input, None)
    }

    fn backward(&mut self, grad_output: &F32Tensor) -> Result<F32Tensor, RusTorchError> {
        // 簡素化されたバックワード実装
        let mut grad = self.output_projection.backward(grad_output)?;

        for block in self.transformer_blocks.iter_mut().rev() {
            grad = block.backward(&grad)?;
        }

        let grad = self.embedding.backward(&grad)?;
        Ok(grad)
    }

    fn parameters(&self) -> Vec<&F32Tensor> {
        let mut params = Vec::new();
        params.extend(self.embedding.parameters());
        for block in &self.transformer_blocks {
            params.extend(block.parameters());
        }
        params.extend(self.output_projection.parameters());
        params
    }

    fn update_parameters(&mut self, learning_rate: f32) -> Result<(), RusTorchError> {
        self.embedding.update_parameters(learning_rate)?;

        for block in &mut self.transformer_blocks {
            block.update_parameters(learning_rate)?;
        }

        self.output_projection.update_parameters(learning_rate)?;
        Ok(())
    }
}