rustytorch/crates/integration/rtx-candle/src/tensor_bridge.rs

//! Tensor bridge between rtx-tensor and Candle tensors

use crate::backend::CandleDevice;
use crate::error::{CandleError, Result};
use candle_core::{DType, Tensor as CandleTensor};
use rtx_tensor::{Device, Tensor};
use tracing::debug;

/// Convert an rtx-tensor Tensor to Candle tensor
pub fn rtx_to_candle(tensor: &Tensor, device: &CandleDevice) -> Result<CandleTensor> {
    let shape: Vec<usize> = tensor.shape().to_vec();
    let data = tensor
        .to_vec_f32()
        .map_err(|e| CandleError::tensor_conversion(format!("Failed to extract f32 data: {e}")))?;

    debug!("Converting rtx tensor {:?} to Candle", shape);

    let candle_device = device.to_candle()?;
    let candle_tensor = CandleTensor::from_vec(data, shape.as_slice(), &candle_device)?;

    Ok(candle_tensor)
}

/// Convert a Candle tensor to rtx-tensor
pub fn candle_to_rtx(tensor: &CandleTensor, device: &Device) -> Result<Tensor> {
    let shape: Vec<usize> = tensor.dims().to_vec();

    // Get data as f32
    let data = match tensor.dtype() {
        DType::F32 => tensor
            .to_vec1::<f32>()
            .or_else(|_| tensor.flatten_all()?.to_vec1::<f32>())?,
        DType::F64 => {
            let f64_data: Vec<f64> = tensor
                .to_vec1()
                .or_else(|_| tensor.flatten_all()?.to_vec1())?;
            f64_data.into_iter().map(|x| x as f32).collect()
        }
        DType::F16 => {
            let f16_data: Vec<half::f16> = tensor
                .to_vec1()
                .or_else(|_| tensor.flatten_all()?.to_vec1())?;
            f16_data.into_iter().map(half::f16::to_f32).collect()
        }
        DType::BF16 => {
            let bf16_data: Vec<half::bf16> = tensor
                .to_vec1()
                .or_else(|_| tensor.flatten_all()?.to_vec1())?;
            bf16_data.into_iter().map(half::bf16::to_f32).collect()
        }
        _ => {
            return Err(CandleError::UnsupportedDType(format!(
                "{:?}",
                tensor.dtype()
            )));
        }
    };

    debug!("Converting Candle tensor {:?} to rtx", shape);

    Tensor::from_vec(data, &shape, device)
        .map_err(|e| CandleError::tensor_conversion(format!("Failed to create tensor: {e}")))
}

/// Convert Candle `DType` to string
pub fn dtype_to_string(dtype: DType) -> &'static str {
    match dtype {
        DType::F32 => "f32",
        DType::F64 => "f64",
        DType::F16 => "f16",
        DType::BF16 => "bf16",
        DType::I64 => "i64",
        DType::I32 => "i32",
        DType::I16 => "i16",
        DType::U32 => "u32",
        DType::U8 => "u8",
        DType::F8E4M3 => "f8e4m3",
        DType::F6E2M3 => "f6e2m3",
        DType::F6E3M2 => "f6e3m2",
        DType::F4 => "f4",
        DType::F8E8M0 => "f8e8m0",
        _ => "unknown",
    }
}

/// Shape utilities
pub struct ShapeUtils;

impl ShapeUtils {
    /// Calculate total number of elements
    pub fn numel(shape: &[usize]) -> usize {
        shape.iter().product()
    }

    /// Check if shapes are equal
    pub fn shapes_equal(a: &[usize], b: &[usize]) -> bool {
        a == b
    }

    /// Validate shape for tensor creation
    pub fn validate_shape(shape: &[usize], data_len: usize) -> bool {
        Self::numel(shape) == data_len
    }
}

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

    #[test]
    fn test_shape_numel() {
        assert_eq!(ShapeUtils::numel(&[2, 3, 4]), 24);
        assert_eq!(ShapeUtils::numel(&[1]), 1);
        assert_eq!(ShapeUtils::numel(&[10, 10]), 100);
    }

    #[test]
    fn test_shapes_equal() {
        assert!(ShapeUtils::shapes_equal(&[2, 3], &[2, 3]));
        assert!(!ShapeUtils::shapes_equal(&[2, 3], &[3, 2]));
    }

    #[test]
    fn test_validate_shape() {
        assert!(ShapeUtils::validate_shape(&[2, 3], 6));
        assert!(!ShapeUtils::validate_shape(&[2, 3], 5));
    }

    #[test]
    fn test_dtype_to_string() {
        assert_eq!(dtype_to_string(DType::F32), "f32");
        assert_eq!(dtype_to_string(DType::F16), "f16");
        assert_eq!(dtype_to_string(DType::BF16), "bf16");
    }
}