amberjs 1.16.1

// 张量操作优化器
//
// 提供多维张量操作、自动微分优化和分布式张量计算能力
// 专为 AI 工作负载设计，支持各种张量运算和梯度计算

/// 张量形状
#[derive(Debug, Clone, PartialEq)]
pub struct TensorShape {
    dims: Vec<usize>,
}
impl TensorShape {
    /// 创建张量形状
    pub fn new(dims: Vec<usize>) -> Self {
        TensorShape { dims }
    }
    /// 获取维度数量
    pub fn rank(&self) -> usize {
        self.dims.len()
    }
    /// 获取总元素数
    pub fn num_elements(&self) -> usize {
        self.dims.iter().product()
    }
    /// 获取指定维度的大小
    pub fn dim(&self, index: usize) -> usize {
        self.dims[index]
    }
    /// 获取所有维度
    pub fn dims(&self) -> &[usize] {
        &self.dims
    }
}
/// 张量数据
#[derive(Debug, Clone)]
pub enum TensorData {
    /// 32 位浮点数
    F32(Vec<f32>),
    /// 64 位浮点数
    F64(Vec<f64>),
    /// 32 位整数
    I32(Vec<i32>),
    /// 64 位整数
    I64(Vec<i64>),
}
impl TensorData {
    /// 获取元素数量
    pub fn len(&self) -> usize {
        match self {
            TensorData::F32(v) => v.len(),
            TensorData::F64(v) => v.len(),
            TensorData::I32(v) => v.len(),
            TensorData::I64(v) => v.len(),
        }
    }
    /// 检查是否为空
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
    /// 获取数据类型
    pub fn data_type(&self) -> &str {
        match self {
            TensorData::F32(_) => "f32",
            TensorData::F64(_) => "f64",
            TensorData::I32(_) => "i32",
            TensorData::I64(_) => "i64",
        }
    }
}
/// 张量结构体
#[derive(Debug, Clone)]
pub struct Tensor {
    shape: TensorShape,
    data: TensorData,
    requires_grad: bool,
}
impl Tensor {
    /// 创建张量
    pub fn new(shape: TensorShape, data: TensorData) -> Self {
        assert_eq!(shape.num_elements(), data.len());
        Tensor {
            shape,
            data,
            requires_grad: false,
        }
    }
    /// 创建需要梯度的张量
    pub fn with_grad(shape: TensorShape, data: TensorData) -> Self {
        assert_eq!(shape.num_elements(), data.len());
        Tensor {
            shape,
            data,
            requires_grad: true,
        }
    }
    /// 获取张量形状
    pub fn shape(&self) -> &TensorShape {
        &self.shape
    }
    /// 获取张量数据
    pub fn data(&self) -> &TensorData {
        &self.data
    }
    /// 检查是否需要梯度
    pub fn requires_grad(&self) -> bool {
        self.requires_grad
    }
    /// 设置需要梯度
    pub fn set_requires_grad(&mut self, requires_grad: bool) {
        self.requires_grad = requires_grad;
    }
    /// 获取 f32 数据引用
    pub fn f32_data(&self) -> Option<&[f32]> {
        match &self.data {
            TensorData::F32(v) => Some(v),
            _ => None,
        }
    }
    /// 获取 f64 数据引用
    pub fn f64_data(&self) -> Option<&[f64]> {
        match &self.data {
            TensorData::F64(v) => Some(v),
            _ => None,
        }
    }
    /// 获取 i32 数据引用
    pub fn i32_data(&self) -> Option<&[i32]> {
        match &self.data {
            TensorData::I32(v) => Some(v),
            _ => None,
        }
    }
    /// 获取 i64 数据引用
    pub fn i64_data(&self) -> Option<&[i64]> {
        match &self.data {
            TensorData::I64(v) => Some(v),
            _ => None,
        }
    }
}
/// 梯度结构
#[derive(Debug, Clone)]
pub struct Gradients {
    tensors: HashMap<String, Tensor>,
}
impl Gradients {
    /// 创建空的梯度
    pub fn new() -> Self {
        Gradients {
            tensors: HashMap::new(),
        }
    }
    /// 添加梯度
    pub fn add(&mut self, name: String, tensor: Tensor) {
        self.tensors.insert(name, tensor);
    }
    /// 获取梯度
    pub fn get(&self, name: &str) -> Option<&Tensor> {
        self.tensors.get(name)
    }
    /// 检查是否包含梯度
    pub fn contains(&self, name: &str) -> bool {
        self.tensors.contains_key(name)
    }
    /// 获取所有梯度名称
    pub fn names(&self) -> Vec<String> {
        self.tensors.keys().cloned().collect()
    }
}
/// 张量分片
#[derive(Debug, Clone)]
pub struct TensorShard {
    tensor: Tensor,
    shard_id: usize,
    total_shards: usize,
}
impl TensorShard {
    /// 创建张量分片
    pub fn new(tensor: Tensor, shard_id: usize, total_shards: usize) -> Self {
        assert!(shard_id < total_shards);
        TensorShard {
            tensor,
            shard_id,
            total_shards,
        }
    }
    /// 获取张量
    pub fn tensor(&self) -> &Tensor {
        &self.tensor
    }
    /// 获取分片 ID
    pub fn shard_id(&self) -> usize {
        self.shard_id
    }
    /// 获取总分片数
    pub fn total_shards(&self) -> usize {
        self.total_shards
    }
}
/// 张量优化器统计信息
#[derive(Debug, Clone)]
pub struct TensorOptimizerStats {
    pub total_operations: u64,
    pub gradient_computations: u64,
    pub distributed_operations: u64,
    pub cache_hits: u64,
    pub cache_misses: u64,
}
/// 张量操作优化器
pub struct TensorOptimizer {
    stats: TensorOptimizerStats,
    operation_count: u64,
    gradient_count: u64,
    distributed_count: u64,
    cache_hits: u64,
    cache_misses: u64,
}
impl TensorOptimizer {
    /// 创建新的张量优化器
    pub fn new() -> Self {
        TensorOptimizer {
            stats: TensorOptimizerStats {
                total_operations: 0,
                gradient_computations: 0,
                distributed_operations: 0,
                cache_hits: 0,
                cache_misses: 0,
            },
            operation_count: 0,
            gradient_count: 0,
            distributed_count: 0,
            cache_hits: 0,
            cache_misses: 0,
        }
    }
    /// 张量矩阵乘法
    pub fn tensor_matmul(&mut self, a: &Tensor, b: &Tensor) -> Tensor {
        assert_eq!(a.shape().rank(), 2);
        assert_eq!(b.shape().rank(), 2);
        assert_eq!(a.shape().dim(1), b.shape().dim(0));
        let (rows, cols) = (a.shape().dim(0), b.shape().dim(1));
        let k: _ = a.shape().dim(1);
        let a_data: _ = a.f32_data().expect("仅支持 f32 数据");
        let b_data: _ = b.f32_data().expect("仅支持 f32 数据");
        let mut result_data = vec![0.0f32; rows * cols];
        for i in 0..rows {
            for j in 0..cols {
                let mut sum = 0.0f32;
                for l in 0..k {
                    sum += a_data[i * k + l] * b_data[l * cols + j];
                }
                result_data[i * cols + j] = sum;
            }
        }
        self.operation_count += 1;
        Tensor::new(
            TensorShape::new(vec![rows, cols]),
            TensorData::F32(result_data),
        )
    }
    /// 计算梯度
    pub fn compute_gradients(&mut self, loss: &Tensor) -> Gradients {
        let mut gradients = Gradients::new();
        // 模拟梯度计算
        match loss.data() {
            TensorData::F32(data) => {
                let grad_data: Vec<f32> = data.iter().map(|&x| 1.0).collect();
                gradients.add(
                    "loss".to_string(),
                    Tensor::new(loss.shape().clone(), TensorData::F32(grad_data)),
                );
            }
            _ => {
                // 其他数据类型的梯度计算
                gradients.add("loss".to_string(), loss.clone());
            }
        }
        self.gradient_count += 1;
        gradients
    }
    /// 分布式张量乘法
    pub fn distributed_matmul(&mut self, shards: &[TensorShard]) -> Tensor {
        assert!(!shards.is_empty());
        let shard_0: _ = &shards[0];
        let rows: _ = shard_0.tensor().shape().dim(0);
        let cols: _ = shard_0.tensor().shape().dim(1);
        // 合并所有分片的结果
        let mut combined_data = vec![0.0f32; rows * cols];
        for shard in shards {
            let tensor: _ = shard.tensor();
            if let Some(data) = tensor.f32_data() {
                for (i, &val) in data.iter().enumerate() {
                    combined_data[i] += val;
                }
            }
        }
        self.distributed_count += shards.len() as u64;
        Tensor::new(
            TensorShape::new(vec![rows, cols]),
            TensorData::F32(combined_data),
        )
    }
    /// 张量加法
    pub fn tensor_add(&mut self, a: &Tensor, b: &Tensor) -> Tensor {
        assert_eq!(a.shape(), b.shape());
        let result_data: _ = match (&a.data, &b.data) {
            (TensorData::F32(a_data), TensorData::F32(b_data)) => TensorData::F32(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x + y)
                    .collect(),
            ),
            (TensorData::F64(a_data), TensorData::F64(b_data)) => TensorData::F64(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x + y)
                    .collect(),
            ),
            (TensorData::I32(a_data), TensorData::I32(b_data)) => TensorData::I32(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x + y)
                    .collect(),
            ),
            (TensorData::I64(a_data), TensorData::I64(b_data)) => TensorData::I64(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x + y)
                    .collect(),
            ),
            _ => panic!("不支持的数据类型组合"),
        };
        self.operation_count += 1;
        Tensor::new(a.shape().clone(), result_data)
    }
    /// 张量减法
    pub fn tensor_sub(&mut self, a: &Tensor, b: &Tensor) -> Tensor {
        assert_eq!(a.shape(), b.shape());
        let result_data: _ = match (&a.data, &b.data) {
            (TensorData::F32(a_data), TensorData::F32(b_data)) => TensorData::F32(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x - y)
                    .collect(),
            ),
            (TensorData::F64(a_data), TensorData::F64(b_data)) => TensorData::F64(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x - y)
                    .collect(),
            ),
            (TensorData::I32(a_data), TensorData::I32(b_data)) => TensorData::I32(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x - y)
                    .collect(),
            ),
            (TensorData::I64(a_data), TensorData::I64(b_data)) => TensorData::I64(
                a_data
                    .iter()
                    .zip(b_data.iter())
                    .map(|(x, y)| x - y)
                    .collect(),
            ),
            _ => panic!("不支持的数据类型组合"),
        };
        self.operation_count += 1;
        Tensor::new(a.shape().clone(), result_data)
    }
    /// 张量标量乘法
    pub fn tensor_scalar_mul(&mut self, tensor: &Tensor, scalar: f32) -> Tensor {
        let result_data: _ = match &tensor.data {
            TensorData::F32(data) => TensorData::F32(data.iter().map(|&x| x * scalar).collect()),
            TensorData::F64(data) => {
                TensorData::F64(data.iter().map(|&x| (x as f32 * scalar) as f64).collect())
            }
            TensorData::I32(data) => {
                TensorData::I32(data.iter().map(|&x| (x as f32 * scalar) as i32).collect())
            }
            TensorData::I64(data) => {
                TensorData::I64(data.iter().map(|&x| (x as f32 * scalar) as i64).collect())
            }
        };
        self.operation_count += 1;
        Tensor::new(tensor.shape().clone(), result_data)
    }
    /// ReLU 激活函数
    pub fn relu(&mut self, tensor: &Tensor) -> Tensor {
        let result_data: _ = match &tensor.data {
            TensorData::F32(data) => TensorData::F32(data.iter().map(|&x| x.max(0.0)).collect()),
            TensorData::F64(data) => TensorData::F64(data.iter().map(|&x| x.max(0.0)).collect()),
            TensorData::I32(data) => TensorData::I32(data.iter().map(|&x| x.max(0)).collect()),
            TensorData::I64(data) => TensorData::I64(data.iter().map(|&x| x.max(0)).collect()),
        };
        self.operation_count += 1;
        Tensor::new(tensor.shape().clone(), result_data)
    }
    /// Softmax 函数
    pub fn softmax(&mut self, tensor: &Tensor) -> Tensor {
        if tensor.shape().rank() != 2 {
            panic!("Softmax 仅支持 2D 张量");
        }
        let rows: _ = tensor.shape().dim(0);
        let cols: _ = tensor.shape().dim(1);
        let mut result_data = vec![0.0f32; rows * cols];
        for i in 0..rows {
            // 计算每行的最大值
            let mut max_val = f32::NEG_INFINITY;
            for j in 0..cols {
                if let Some(data) = tensor.f32_data() {
                    max_val = max_val.max(data[i * cols + j]);
                }
            }
            // 计算每行的指数和
            let mut exp_sum = 0.0f32;
            for j in 0..cols {
                if let Some(data) = tensor.f32_data() {
                    let exp_val: _ = (data[i * cols + j] - max_val).exp();
                    exp_sum += exp_val;
                    result_data[i * cols + j] = exp_val;
                }
            }
            // 归一化
            for j in 0..cols {
                result_data[i * cols + j] /= exp_sum;
            }
        }
        self.operation_count += 1;
        Tensor::new(tensor.shape().clone(), TensorData::F32(result_data))
    }
    /// 获取统计信息
    pub fn get_stats(&self) -> TensorOptimizerStats {
        TensorOptimizerStats {
            total_operations: self.operation_count,
            gradient_computations: self.gradient_count,
            distributed_operations: self.distributed_count,
            cache_hits: self.cache_hits,
            cache_misses: self.cache_misses,
        }
    }
}
impl Default for TensorOptimizer {
    fn default() -> Self {
        Self::new()
    }
}
use std::collections::{BTreeMap, HashMap};