foxhunt/crates/ml/src/tensor_ops.rs

//!
//! Tensor operations and utilities for ML models
//!
//! Provides optimized tensor operations for high-frequency trading models
//! with focus on ultra-low latency inference.

use candle_core::{Device, Result as CandleResult, Tensor};

// Use Tensor directly for integer operations - no wrapper needed

/// Tensor operation utilities
#[derive(Debug)]
pub struct TensorOps;

impl TensorOps {
    /// Create a new integer tensor from vector
    pub fn from_vec_i32(data: Vec<i32>, device: &Device) -> CandleResult<Tensor> {
        let f32_data: Vec<f32> = data.iter().map(|&x| x as f32).collect();
        Tensor::from_vec(f32_data, (data.len(),), device)
    }

    /// Create a new integer tensor from slice
    pub fn from_slice_i32(data: &[i32], shape: &[usize], device: &Device) -> CandleResult<Tensor> {
        let f32_data: Vec<f32> = data.iter().map(|&x| x as f32).collect();
        Tensor::from_slice(&f32_data, shape, device)
    }

    /// Convert tensor to i32 vector
    pub fn to_vec_i32(tensor: &Tensor) -> CandleResult<Vec<i32>> {
        let f32_vec = tensor.to_vec1::<f32>()?;
        Ok(f32_vec.iter().map(|&x| x as i32).collect())
    }

    /// Apply softmax operation with numerical stability
    pub fn stable_softmax(input: &Tensor, dim: usize) -> CandleResult<Tensor> {
        let max_vals = input.max_keepdim(dim)?;
        let shifted = input.broadcast_sub(&max_vals)?;
        let exp_vals = shifted.exp()?;
        let sum_exp = exp_vals.sum_keepdim(dim)?;
        exp_vals.broadcast_div(&sum_exp)
    }

    /// Clamp tensor values between min and max
    pub fn clamp(input: &Tensor, min_val: f64, max_val: f64) -> CandleResult<Tensor> {
        let min_tensor = Tensor::full(min_val as f32, input.shape(), input.device())?;
        let max_tensor = Tensor::full(max_val as f32, input.shape(), input.device())?;
        input.clamp(&min_tensor, &max_tensor)
    }

    /// Normalize tensor to unit length
    pub fn normalize(input: &Tensor, dim: usize) -> CandleResult<Tensor> {
        let norm = input.sqr()?.sum_keepdim(dim)?.sqrt()?;
        let epsilon = Tensor::full(1e-8_f32, norm.shape(), norm.device())?;
        let norm_safe = norm.add(&epsilon)?;
        input.broadcast_div(&norm_safe)
    }

    /// Negate tensor (equivalent to unary minus operator)
    pub fn negate(input: &Tensor) -> CandleResult<Tensor> {
        let zero = Tensor::zeros(input.shape(), input.dtype(), input.device())?;
        zero.sub(input)
    }

    /// Element-wise minimum between two tensors
    pub fn elementwise_min(a: &Tensor, b: &Tensor) -> CandleResult<Tensor> {
        let diff = a.sub(b)?;
        let mask = diff.lt(&Tensor::zeros(diff.shape(), diff.dtype(), diff.device())?)?;
        let mask_f32 = mask.to_dtype(a.dtype())?;
        let one_minus_mask =
            Tensor::ones(mask_f32.shape(), mask_f32.dtype(), mask_f32.device())?.sub(&mask_f32)?;
        a.mul(&mask_f32)?.add(&b.mul(&one_minus_mask)?)
    }

    /// Multiply tensor by scalar (handles f32/f64 conversion)
    pub fn scalar_mul(tensor: &Tensor, scalar: f64) -> CandleResult<Tensor> {
        let scalar_tensor = Tensor::full(scalar as f32, tensor.shape(), tensor.device())?;
        tensor.mul(&scalar_tensor)
    }
}

/// Extension trait for integer tensor operations
pub trait IntegerTensorExt {
    /// Create new integer tensor from vector
    fn from_vec_i32(data: Vec<i32>, device: &Device) -> CandleResult<Tensor>;

    /// Convert to i32 vector
    fn to_vec_i32(&self) -> CandleResult<Vec<i32>>;
}

impl IntegerTensorExt for Tensor {
    fn from_vec_i32(data: Vec<i32>, device: &Device) -> CandleResult<Self> {
        TensorOps::from_vec_i32(data, device)
    }

    fn to_vec_i32(&self) -> CandleResult<Vec<i32>> {
        TensorOps::to_vec_i32(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::Device;

    #[test]
    fn test_integer_tensor_creation() -> CandleResult<()> {
        let device = Device::Cpu;
        let data = vec![1, 2, 3, 4, 5];
        // Use IntegerTensorExt trait method
        let tensor = Tensor::from_vec_i32(data.clone(), &device)?;
        let result = tensor.to_vec_i32()?;
        assert_eq!(data, result);
        Ok(())
    }

    #[test]
    fn test_stable_softmax() -> CandleResult<()> {
        let device = Device::Cpu;
        let data = vec![1.0_f32, 2.0, 3.0];
        let tensor = Tensor::from_vec(data, 3, &device)?;
        let softmax = TensorOps::stable_softmax(&tensor, 0)?;
        let result: Vec<f32> = softmax.to_vec1()?;

        // Check that probabilities sum to 1
        let sum: f32 = result.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
        Ok(())
    }

    #[test]
    fn test_clamp() -> CandleResult<()> {
        let device = Device::Cpu;
        let data = vec![-2.0_f32, -1.0, 0.0, 1.0, 2.0];
        let tensor = Tensor::from_vec(data, 5, &device)?;
        let clamped = TensorOps::clamp(&tensor, -1.0, 1.0)?;
        let result: Vec<f32> = clamped.to_vec1()?;

        for val in result {
            assert!(val >= -1.0 && val <= 1.0);
        }
        Ok(())
    }
}