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

//! Tensor bridge between rtx-tensor and Burn tensors
//!
//! Provides conversion functions for zero-copy (when possible) tensor sharing.

use crate::error::{BurnError, Result};
use rtx_tensor::{DType, Device, Tensor};

/// Convert an rtx-tensor Tensor to Burn tensor data
///
/// Returns the data as a Vec<f32> along with the shape for Burn consumption.
pub fn rtx_to_burn(tensor: &Tensor) -> Result<(Vec<f32>, Vec<usize>)> {
    let shape = tensor.shape().to_vec();

    let data = tensor
        .to_vec_f32()
        .map_err(|e| BurnError::tensor_conversion(format!("Failed to extract f32 data: {e}")))?;

    Ok((data, shape))
}

/// Convert Burn tensor data back to rtx-tensor
///
/// Takes the data and shape from a Burn tensor and creates an rtx Tensor.
pub fn burn_to_rtx(data: Vec<f32>, shape: Vec<usize>, device: &Device) -> Result<Tensor> {
    Tensor::from_vec(data, &shape, device)
        .map_err(|e| BurnError::tensor_conversion(format!("Failed to create tensor: {e}")))
}

/// Tensor shape utilities
pub struct ShapeUtils;

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

    /// Check if two shapes are compatible for broadcasting
    pub fn broadcast_compatible(a: &[usize], b: &[usize]) -> bool {
        let max_dims = a.len().max(b.len());

        for i in 0..max_dims {
            let dim_a = if i < a.len() { a[a.len() - 1 - i] } else { 1 };
            let dim_b = if i < b.len() { b[b.len() - 1 - i] } else { 1 };

            if dim_a != dim_b && dim_a != 1 && dim_b != 1 {
                return false;
            }
        }
        true
    }

    /// Calculate broadcast result shape
    pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Vec<usize>> {
        if !Self::broadcast_compatible(a, b) {
            return None;
        }

        let max_dims = a.len().max(b.len());
        let mut result = Vec::with_capacity(max_dims);

        for i in 0..max_dims {
            let dim_a = if i < a.len() { a[a.len() - 1 - i] } else { 1 };
            let dim_b = if i < b.len() { b[b.len() - 1 - i] } else { 1 };
            result.push(dim_a.max(dim_b));
        }

        result.reverse();
        Some(result)
    }

    /// Convert shape to Burn format (fixed-size array via const generics workaround)
    pub fn to_dims<const N: usize>(shape: &[usize]) -> Option<[usize; N]> {
        if shape.len() != N {
            return None;
        }
        let mut dims = [0usize; N];
        dims.copy_from_slice(shape);
        Some(dims)
    }
}

/// Data type conversion utilities
pub struct DTypeUtils;

impl DTypeUtils {
    /// Check if dtype is supported by Burn
    pub fn is_supported(dtype: DType) -> bool {
        matches!(dtype, DType::F32 | DType::F64 | DType::I32 | DType::I64)
    }

    /// Get the element size in bytes
    pub fn element_size(dtype: DType) -> usize {
        match dtype {
            DType::F32 | DType::I32 => 4,
            DType::F64 | DType::I64 => 8,
            DType::F16 | DType::BF16 => 2,
            DType::I8 | DType::U8 => 1,
            _ => 4, // Default to 4 bytes
        }
    }
}

#[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, 1]), 1);
        assert_eq!(ShapeUtils::numel(&[10]), 10);
    }

    #[test]
    fn test_broadcast_compatible() {
        assert!(ShapeUtils::broadcast_compatible(&[3, 4], &[3, 4]));
        assert!(ShapeUtils::broadcast_compatible(&[3, 4], &[1, 4]));
        assert!(ShapeUtils::broadcast_compatible(&[3, 4], &[4]));
        assert!(ShapeUtils::broadcast_compatible(&[3, 1], &[1, 4]));
        assert!(!ShapeUtils::broadcast_compatible(&[3, 4], &[2, 4]));
    }

    #[test]
    fn test_broadcast_shape() {
        assert_eq!(
            ShapeUtils::broadcast_shape(&[3, 1], &[1, 4]),
            Some(vec![3, 4])
        );
        assert_eq!(
            ShapeUtils::broadcast_shape(&[2, 3, 4], &[4]),
            Some(vec![2, 3, 4])
        );
        assert_eq!(ShapeUtils::broadcast_shape(&[3, 4], &[2, 4]), None);
    }

    #[test]
    fn test_to_dims() {
        let shape = vec![2, 3, 4];
        let dims: Option<[usize; 3]> = ShapeUtils::to_dims(&shape);
        assert_eq!(dims, Some([2, 3, 4]));

        let dims: Option<[usize; 2]> = ShapeUtils::to_dims(&shape);
        assert_eq!(dims, None);
    }

    #[test]
    fn test_dtype_supported() {
        assert!(DTypeUtils::is_supported(DType::F32));
        assert!(DTypeUtils::is_supported(DType::I64));
    }

    #[test]
    fn test_element_size() {
        assert_eq!(DTypeUtils::element_size(DType::F32), 4);
        assert_eq!(DTypeUtils::element_size(DType::F64), 8);
        assert_eq!(DTypeUtils::element_size(DType::F16), 2);
    }
}