rustytorch/crates/training/rtx-rl/src/tensor_utils.rs

use crate::{Result, RLError};
use rtx_tensor::{Tensor, Device, DType, Shape};
use rand::{thread_rng, Rng};

/// Utility functions for tensor operations not yet available in rtx-tensor
pub struct TensorUtils;

impl TensorUtils {
    /// Create a tensor filled with random normal values
    pub fn randn(shape: &[usize], dtype: DType, device: &Device) -> Result<Tensor> {
        let shape_obj = Shape::new(shape.to_vec()).map_err(|e| RLError::TensorError(e.to_string()))?;
        let numel = shape_obj.numel();

        let mut rng = thread_rng();
        let data: Vec<f32> = (0..numel)
        .map(|_| rng.gen_range(-1.0..1.0)) // Simple uniform random for now
        .collect();

        Tensor::from_data(data, shape_obj, device).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Create tensor from vec with proper shape
    pub fn from_vec(data: Vec<f32>, shape: &[usize], device: &Device) -> Result<Tensor> {
        let shape_obj = Shape::new(shape.to_vec()).map_err(|e| RLError::TensorError(e.to_string()))?;
        Tensor::from_data(data, shape_obj, device).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Stack tensors along a dimension
    pub fn stack(tensors: &[Tensor], dim: usize) -> Result<Tensor> {
        if tensors.is_empty() {
            return Err(RLError::TensorError("Cannot stack empty tensor list".to_string()));
        }

        // For now, just return the first tensor as a placeholder
        // In real implementation, we'd properly stack the tensors
        Ok(tensors[0].clone())
    }

    /// Convert tensor to vec (simplified)
    pub async fn to_vec(tensor: &Tensor) -> Result<Vec<f32>> {
        tensor.to_cpu().map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Get scalar value from tensor
    pub async fn item(tensor: &Tensor) -> Result<f32> {
        let data = tensor.to_cpu().map_err(|e| RLError::TensorError(e.to_string()))?;
        if data.is_empty() {
            return Err(RLError::TensorError("Empty tensor has no item".to_string()));
        }
        Ok(data[0])
    }

    /// Compute tensor mean (simplified)
    pub fn mean(tensor: &Tensor) -> Result<Tensor> {
        // Simplified implementation - sum and divide by numel
        let sum = tensor.sum(None).map_err(|e| RLError::TensorError(e.to_string()))?;
        let numel = tensor.numel() as f32;
        sum.scalar_mul(1.0 / numel).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Element-wise exponential
    pub fn exp(tensor: &Tensor) -> Result<Tensor> {
        // Simplified - return input for now (would need proper exp implementation)
        Ok(tensor.clone())
    }

    /// Clamp tensor values
    pub fn clamp(tensor: &Tensor, min: f32, max: f32) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Tanh activation
    pub fn tanh(tensor: &Tensor) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Softmax function
    pub fn softmax(tensor: &Tensor, dim: i32) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Log-softmax function
    pub fn log_softmax(tensor: &Tensor, dim: i32) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Power function
    pub fn pow(tensor: &Tensor, exp: f32) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Sigmoid function
    pub fn sigmoid(tensor: &Tensor) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Min of two tensors
    pub fn min(tensor1: &Tensor, tensor2: &Tensor) -> Result<Tensor> {
        tensor1.minimum(tensor2).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Element-wise comparison (greater than)
    pub fn gt(tensor1: &Tensor, tensor2: &Tensor) -> Result<Tensor> {
        // Simplified - return ones for now
        Tensor::ones(tensor1.shape().clone(), tensor1.device()).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Convert to dtype
    pub fn to_dtype(tensor: &Tensor, dtype: DType) -> Result<Tensor> {
        // Simplified - return input for now
        Ok(tensor.clone())
    }

    /// Unsqueeze (add dimension)
    pub fn unsqueeze(tensor: &Tensor, dim: usize) -> Result<Tensor> {
        // Simplified implementation
        let mut new_shape = tensor.shape().dims().to_vec();
        new_shape.insert(dim, 1);
        tensor.view(new_shape).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Squeeze (remove dimension)
    pub fn squeeze(tensor: &Tensor, dim: usize) -> Result<Tensor> {
        // Simplified implementation
        let mut new_shape = tensor.shape().dims().to_vec();
        if dim < new_shape.len() && new_shape[dim] == 1 {
            new_shape.remove(dim);
        }
        tensor.view(new_shape).map_err(|e| RLError::TensorError(e.to_string()))
    }

    /// Detach tensor from computation graph
    pub fn detach(tensor: &Tensor) -> Result<Tensor> {
        // Simplified - return clone for now
        Ok(tensor.clone())
    }
}