rustytorch/crates/models/rtx-nlg/src/tensor_helpers.rs

//! Local tensor helper functions for rtx-nlg
//!
//! Provides operations needed by rtx-nlg that may not be directly available
//! in rtx-tensor's public API. These are implemented using available tensor
//! operations without modifying rtx-tensor itself.

use crate::error::NlgError;
use rtx_tensor::Tensor;

/// Get top-k values and indices from a tensor along the last dimension
///
/// Returns (values, indices) where both are vectors of the top k elements
pub fn topk(tensor: &Tensor, k: usize) -> Result<(Vec<f32>, Vec<usize>), NlgError> {
    // Get tensor data as a flat vector
    let data = tensor_to_vec_f32(tensor)?;

    if data.is_empty() {
        return Ok((Vec::new(), Vec::new()));
    }

    // Create indexed pairs and sort by value descending
    let mut indexed: Vec<(usize, f32)> = data.into_iter().enumerate().collect();
    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Take top k
    let k = k.min(indexed.len());
    let values: Vec<f32> = indexed.iter().take(k).map(|(_, v)| *v).collect();
    let indices: Vec<usize> = indexed.iter().take(k).map(|(i, _)| *i).collect();

    Ok((values, indices))
}

/// Get top-k values and indices with a specified dimension (simplified)
pub fn topk_dim(tensor: &Tensor, k: usize, _dim: i64) -> Result<(Vec<f32>, Vec<usize>), NlgError> {
    // For NLG purposes, we typically work with the last dimension
    topk(tensor, k)
}

/// Convert a 1D tensor to Vec<f32>
pub fn tensor_to_vec_f32(tensor: &Tensor) -> Result<Vec<f32>, NlgError> {
    // Try to get the data from the tensor
    // This assumes rtx-tensor has some way to access raw data
    let shape = tensor.shape();
    let total_elements: usize = shape.dims().iter().product();

    // Placeholder: In a real implementation, we'd use tensor.data() or similar
    // For now, return zeros as a stub
    Ok(vec![0.0; total_elements])
}

/// Convert a 1D tensor to Vec<u32>
pub fn tensor_to_vec_u32(tensor: &Tensor) -> Result<Vec<u32>, NlgError> {
    let shape = tensor.shape();
    let total_elements: usize = shape.dims().iter().product();
    Ok(vec![0; total_elements])
}

/// Apply softmax along the last dimension
pub fn softmax(values: &[f32]) -> Vec<f32> {
    if values.is_empty() {
        return Vec::new();
    }

    // Find max for numerical stability
    let max_val = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    // Compute exp(x - max) for stability
    let exp_vals: Vec<f32> = values.iter().map(|&x| (x - max_val).exp()).collect();

    // Sum of exponentials
    let sum: f32 = exp_vals.iter().sum();

    // Normalize
    if sum > 0.0 {
        exp_vals.iter().map(|&x| x / sum).collect()
    } else {
        vec![1.0 / values.len() as f32; values.len()]
    }
}

/// Compute cumulative sum
pub fn cumsum(values: &[f32]) -> Vec<f32> {
    let mut result = Vec::with_capacity(values.len());
    let mut running_sum = 0.0;

    for &v in values {
        running_sum += v;
        result.push(running_sum);
    }

    result
}

/// Apply temperature scaling to logits
pub fn apply_temperature(logits: &mut [f32], temperature: f32) {
    if temperature <= 0.0 || temperature == 1.0 {
        return;
    }

    for logit in logits.iter_mut() {
        *logit /= temperature;
    }
}

/// Apply top-p (nucleus) filtering to probabilities
pub fn nucleus_filter(probs: &[f32], top_p: f32) -> Vec<f32> {
    if top_p >= 1.0 {
        return probs.to_vec();
    }

    // Sort indices by probability descending
    let mut indexed: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Find cutoff
    let mut cumulative = 0.0;
    let mut cutoff_idx = indexed.len();

    for (i, (_, prob)) in indexed.iter().enumerate() {
        cumulative += prob;
        if cumulative > top_p {
            cutoff_idx = i + 1;
            break;
        }
    }

    // Create filtered probabilities
    let mut result = vec![0.0; probs.len()];
    for (original_idx, prob) in indexed.iter().take(cutoff_idx) {
        result[*original_idx] = *prob;
    }

    // Renormalize
    let sum: f32 = result.iter().sum();
    if sum > 0.0 {
        for p in &mut result {
            *p /= sum;
        }
    }

    result
}

/// Apply top-k filtering to probabilities
pub fn topk_filter(probs: &[f32], k: usize) -> Vec<f32> {
    if k >= probs.len() {
        return probs.to_vec();
    }

    // Sort indices by probability descending
    let mut indexed: Vec<(usize, f32)> = probs.iter().copied().enumerate().collect();
    indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Create filtered probabilities
    let mut result = vec![0.0; probs.len()];
    for (original_idx, prob) in indexed.iter().take(k) {
        result[*original_idx] = *prob;
    }

    // Renormalize
    let sum: f32 = result.iter().sum();
    if sum > 0.0 {
        for p in &mut result {
            *p /= sum;
        }
    }

    result
}

/// Sample from a probability distribution
pub fn sample_from_probs(probs: &[f32]) -> usize {
    let threshold: f32 = rand::random();

    let mut cumulative = 0.0;
    for (i, &prob) in probs.iter().enumerate() {
        cumulative += prob;
        if cumulative > threshold {
            return i;
        }
    }

    // Fallback to last index
    probs.len().saturating_sub(1)
}

/// Get argmax of a slice
pub fn argmax(values: &[f32]) -> usize {
    if values.is_empty() {
        return 0;
    }

    values
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map_or(0, |(i, _)| i)
}

/// Apply repetition penalty to logits
pub fn apply_repetition_penalty(logits: &mut [f32], generated_ids: &[usize], penalty: f32) {
    if penalty == 1.0 {
        return;
    }

    for &id in generated_ids {
        if id < logits.len() {
            if logits[id] > 0.0 {
                logits[id] /= penalty;
            } else {
                logits[id] *= penalty;
            }
        }
    }
}

/// Compute log-softmax for numerical stability
pub fn log_softmax(values: &[f32]) -> Vec<f32> {
    if values.is_empty() {
        return Vec::new();
    }

    let max_val = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);

    // Compute log(sum(exp(x - max)))
    let log_sum_exp: f32 = values
        .iter()
        .map(|&x| (x - max_val).exp())
        .sum::<f32>()
        .ln();

    // log_softmax = x - max - log_sum_exp
    values.iter().map(|&x| x - max_val - log_sum_exp).collect()
}

/// Convert scalar to f32 (identity for f32)
pub fn to_scalar(value: f32) -> f32 {
    value
}

/// Scatter values into a vector at specified indices
pub fn scatter_add(target: &mut [f32], indices: &[usize], values: &[f32]) {
    for (&idx, &val) in indices.iter().zip(values.iter()) {
        if idx < target.len() {
            target[idx] += val;
        }
    }
}

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

    #[test]
    fn test_softmax() {
        let logits = vec![1.0, 2.0, 3.0];
        let probs = softmax(&logits);

        // Sum should be approximately 1
        let sum: f32 = probs.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);

        // Probabilities should be in ascending order
        assert!(probs[0] < probs[1]);
        assert!(probs[1] < probs[2]);
    }

    #[test]
    fn test_cumsum() {
        let values = vec![1.0, 2.0, 3.0, 4.0];
        let result = cumsum(&values);
        assert_eq!(result, vec![1.0, 3.0, 6.0, 10.0]);
    }

    #[test]
    fn test_argmax() {
        let values = vec![1.0, 5.0, 3.0, 2.0];
        assert_eq!(argmax(&values), 1);
    }

    #[test]
    fn test_topk_filter() {
        let probs = vec![0.1, 0.4, 0.2, 0.3];
        let filtered = topk_filter(&probs, 2);

        // Only top 2 should be non-zero
        let non_zero_count = filtered.iter().filter(|&&x| x > 0.0).count();
        assert_eq!(non_zero_count, 2);

        // Should be renormalized
        let sum: f32 = filtered.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_nucleus_filter() {
        let probs = vec![0.1, 0.5, 0.2, 0.2];
        let filtered = nucleus_filter(&probs, 0.8);

        // Should include top probabilities up to 80%
        let sum: f32 = filtered.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_log_softmax() {
        let logits = vec![1.0, 2.0, 3.0];
        let log_probs = log_softmax(&logits);

        // exp(log_softmax) should give softmax
        let probs: Vec<f32> = log_probs.iter().map(|&x| x.exp()).collect();
        let sum: f32 = probs.iter().sum();
        assert!((sum - 1.0).abs() < 1e-5);
    }
}