rustytorch/crates/core/rtx-cubecl/src/backend.rs

//! `CubeCL` Backend implementation for `RustyTorch`++.
//!
//! This module provides the `CubeclBackend` type that implements the `Backend` trait
//! from `rtx-backend`, enabling `CubeCL` to be used as a compute backend for tensor operations.
//!
//! The `CubeCL` backend provides portable GPU kernels that compile to:
//! - CUDA PTX (NVIDIA GPUs)
//! - WGSL (WebGPU)
//! - HIP (AMD GPUs)
//! - Metal MSL (via WebGPU)
//! - SPIR-V (Vulkan)
//!
//! ## Tiered Kernel Strategy
//!
//! For optimal performance, `RustyTorch`++ uses a tiered approach:
//!
//! - **`CubeCL` (Tier 1)**: Portable ops (~60% of operations)
//!   - Elementwise: add, sub, mul, div, exp, log, sqrt
//!   - Activations: relu, gelu, sigmoid, tanh, silu
//!   - Reductions: sum, mean, max, min
//!   - Shape ops: reshape, transpose
//!
//! - **Native (Tier 2)**: Performance-critical ops (~40% of operations)
//!   - Large matmuls → cuBLAS/rocBLAS
//!   - `FlashAttention` → hand-optimized CUDA
//!   - Convolutions → cuDNN

use std::fmt::Debug;

use rtx_backend::{Backend, BoolU8, DeviceId, DeviceOps};

use crate::client;
use crate::device::CubeclDevice;
use crate::ops;
use crate::runtime::RuntimeBackend;
use crate::tensor::CubeclTensorPrimitive;

/// CubeCL-based backend for portable GPU operations.
///
/// This backend uses `CubeCL` to compile GPU kernels from Rust code,
/// enabling the same kernel source to run on CUDA, WebGPU, HIP, etc.
///
/// # Example
///
/// ```rust,ignore
/// use rtx_cubecl::CubeclBackend;
/// use rtx_backend::Backend;
///
/// type MyBackend = CubeclBackend;
///
/// let device = CubeclBackend::default_device();
/// let tensor = MyBackend::zeros::<2>([3, 4], &device);
/// ```
#[derive(Clone, Debug, Default)]
pub struct CubeclBackend;

/// Device operations for `CubeCL`.
impl DeviceOps<CubeclBackend> for CubeclDevice {
    fn id(&self) -> DeviceId {
        match self.backend {
            RuntimeBackend::Cuda => DeviceId::Cuda(self.index),
            RuntimeBackend::Hip => DeviceId::Rocm(self.index),
            RuntimeBackend::Wgpu => DeviceId::WebGpu(self.index),
            RuntimeBackend::Cpu => DeviceId::Cpu,
        }
    }

    fn memory_capacity(&self) -> usize {
        // TODO: Query actual device memory capacity
        match self.backend {
            RuntimeBackend::Cpu => {
                // Return system memory as approximation
                16 * 1024 * 1024 * 1024 // 16 GB default
            }
            RuntimeBackend::Cuda | RuntimeBackend::Hip | RuntimeBackend::Wgpu => {
                // Default GPU memory
                8 * 1024 * 1024 * 1024 // 8 GB default
            }
        }
    }

    fn memory_available(&self) -> usize {
        // TODO: Query actual available memory
        self.memory_capacity() / 2 // Conservative estimate
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        match self.backend {
            RuntimeBackend::Cuda => Some((8, 0)), // Default to Ampere
            RuntimeBackend::Hip => Some((9, 0)),  // Default to RDNA 2
            RuntimeBackend::Wgpu => None,         // WebGPU doesn't have compute capability
            RuntimeBackend::Cpu => None,
        }
    }

    fn synchronize(&self) {
        // TODO: Implement actual synchronization
        // For now, this is a no-op as we don't have actual compute yet
    }

    fn is_available(&self) -> bool {
        match self.backend {
            RuntimeBackend::Cpu => true,
            #[cfg(feature = "cuda")]
            RuntimeBackend::Cuda => true, // TODO: Check actual CUDA availability
            #[cfg(not(feature = "cuda"))]
            RuntimeBackend::Cuda => false,
            #[cfg(feature = "hip")]
            RuntimeBackend::Hip => true, // TODO: Check actual HIP availability
            #[cfg(not(feature = "hip"))]
            RuntimeBackend::Hip => false,
            #[cfg(feature = "wgpu")]
            RuntimeBackend::Wgpu => true, // TODO: Check actual WebGPU availability
            #[cfg(not(feature = "wgpu"))]
            RuntimeBackend::Wgpu => false,
        }
    }
}

impl Backend for CubeclBackend {
    type TensorPrimitive<const D: usize> = CubeclTensorPrimitive<D>;
    type Device = CubeclDevice;
    type FloatElem = f32;
    type IntElem = i32;
    type BoolElem = BoolU8;

    fn name() -> &'static str {
        "CubeCL"
    }

    fn seed(seed: u64) {
        // TODO: Implement RNG seeding for CubeCL
        let _ = seed;
    }

    // ==================== Tensor Creation ====================

    fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        let numel: usize = shape.iter().product();
        let data = vec![0.0f32; numel];

        // Try to allocate on GPU, fall back to CPU
        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(&data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        // CPU fallback
        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        let numel: usize = shape.iter().product();
        let data = vec![1.0f32; numel];

        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(&data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    fn full<const D: usize>(
        shape: [usize; D],
        fill_value: Self::FloatElem,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        let numel: usize = shape.iter().product();
        let data = vec![fill_value; numel];

        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(&data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        // Use simple LCG for now, replace with proper RNG later
        let numel: usize = shape.iter().product();
        let mut seed = 12345u64;
        let data: Vec<f32> = (0..numel)
            .map(|_| {
                seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
                ((seed >> 16) & 0x7fff) as f32 / 32767.0
            })
            .collect();

        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(&data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        // Box-Muller transform for normal distribution
        let numel: usize = shape.iter().product();
        let mut seed = 67890u64;
        let data: Vec<f32> = (0..numel)
            .map(|_| {
                // Generate two uniform random numbers
                seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
                let u1 = ((seed >> 16) & 0x7fff) as f32 / 32768.0 + 1e-10;
                seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
                let u2 = ((seed >> 16) & 0x7fff) as f32 / 32768.0;
                // Box-Muller transform
                (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
            })
            .collect();

        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(&data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    fn from_data<const D: usize>(
        data: &[Self::FloatElem],
        shape: [usize; D],
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        if let Ok(client) = client::get_or_create(device)
            && let Ok(handle) = client.create_f32(data)
        {
            return CubeclTensorPrimitive::from_handle(shape, handle, device.clone());
        }

        let bytes: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
        CubeclTensorPrimitive::from_cpu_data(shape, bytes, device.clone())
    }

    // ==================== Binary Operations ====================

    fn add<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_binary(lhs, rhs, |a, b| a + b)
    }

    fn sub<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_binary(lhs, rhs, |a, b| a - b)
    }

    fn mul<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_binary(lhs, rhs, |a, b| a * b)
    }

    fn div<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_binary(lhs, rhs, |a, b| a / b)
    }

    fn matmul(
        lhs: Self::TensorPrimitive<2>,
        rhs: Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<2> {
        ops::matmul(lhs, rhs)
    }

    fn bmm(
        lhs: Self::TensorPrimitive<3>,
        rhs: Self::TensorPrimitive<3>,
    ) -> Self::TensorPrimitive<3> {
        ops::bmm(lhs, rhs)
    }

    // ==================== Unary Operations ====================

    fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, |x| -x)
    }

    fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::exp)
    }

    fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::ln)
    }

    fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::sqrt)
    }

    fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::abs)
    }

    fn sin<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::sin)
    }

    fn cos<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::cos)
    }

    fn pow<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        exp: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, move |x| x.powf(exp))
    }

    fn clamp<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        min: Self::FloatElem,
        max: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, move |x| x.clamp(min, max))
    }

    // ==================== Activation Functions ====================

    fn relu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, |x| x.max(0.0))
    }

    fn sigmoid<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, |x| 1.0 / (1.0 + (-x).exp()))
    }

    fn tanh<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(tensor, f32::tanh)
    }

    fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        // GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
        ops::elementwise_unary(tensor, |x| {
            let sqrt_2_over_pi = 0.797_884_6_f32;
            let coef = 0.044715f32;
            let inner = sqrt_2_over_pi * (x + coef * x * x * x);
            0.5 * x * (1.0 + inner.tanh())
        })
    }

    fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        // SiLU (Swish): x * sigmoid(x)
        ops::elementwise_unary(tensor, |x| x * (1.0 / (1.0 + (-x).exp())))
    }

    fn leaky_relu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        negative_slope: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(
            tensor,
            move |x| if x > 0.0 { x } else { negative_slope * x },
        )
    }

    fn elu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        alpha: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::elementwise_unary(
            tensor,
            move |x| if x > 0.0 { x } else { alpha * (x.exp() - 1.0) },
        )
    }

    // ==================== Reduction Operations ====================

    fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::sum(tensor)
    }

    fn sum_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Implement proper dimension reduction
        // For now, return input (placeholder)
        let _ = dim;
        tensor
    }

    fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::mean(tensor)
    }

    fn mean_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Implement proper dimension reduction
        let _ = dim;
        tensor
    }

    fn var<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::var(tensor)
    }

    fn var_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Implement proper dimension reduction
        let _ = dim;
        tensor
    }

    fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::max(tensor)
    }

    fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::min(tensor)
    }

    // ==================== Shape Operations ====================

    fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D] {
        tensor.shape
    }

    fn reshape<const D1: usize, const D2: usize>(
        tensor: Self::TensorPrimitive<D1>,
        shape: [usize; D2],
    ) -> Self::TensorPrimitive<D2> {
        tensor.reshape(shape).expect("Reshape failed")
    }

    fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        // TODO: Launch transpose kernel
        let mut new_shape = tensor.shape;
        if D >= 2 {
            new_shape.swap(D - 2, D - 1);
        }
        CubeclTensorPrimitive::new(new_shape, tensor.device)
    }

    fn swap_dims<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim1: usize,
        dim2: usize,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Launch permute kernel
        let mut new_shape = tensor.shape;
        new_shape.swap(dim1, dim2);
        CubeclTensorPrimitive::new(new_shape, tensor.device)
    }

    // ==================== LLM-Specific Operations ====================
    // These delegate to native kernels for maximum performance

    fn flash_attention(
        query: Self::TensorPrimitive<4>,
        key: Self::TensorPrimitive<4>,
        value: Self::TensorPrimitive<4>,
        mask: Option<&Self::TensorPrimitive<4>>,
        scale: Self::FloatElem,
        causal: bool,
    ) -> Self::TensorPrimitive<4> {
        // TODO: Delegate to native FlashAttention kernel
        // CubeCL is not used here for maximum performance
        let _ = (key, value, mask, scale, causal);
        query
    }

    fn softmax<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::softmax(tensor, dim)
    }

    fn layer_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        bias: Option<&Self::TensorPrimitive<1>>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Launch layer_norm kernel
        let _ = (weight, bias, eps);
        tensor
    }

    fn rms_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Launch rms_norm kernel
        let _ = (weight, eps);
        tensor
    }

    fn rope<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        cos: &Self::TensorPrimitive<2>,
        sin: &Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Launch rope kernel
        let _ = (cos, sin);
        tensor
    }

    // ==================== Comparison Operations ====================

    fn gt_scalar<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        value: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Launch gt_scalar kernel
        let _ = value;
        tensor
    }

    // ==================== Convolution Operations ====================

    fn conv2d(
        input: Self::TensorPrimitive<4>,
        weight: &Self::TensorPrimitive<4>,
        bias: Option<&Self::TensorPrimitive<1>>,
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        groups: usize,
    ) -> Self::TensorPrimitive<4> {
        // TODO: Delegate to cuDNN or implement CubeCL convolution
        let _ = (bias, groups);
        let batch = input.shape[0];
        let out_channels = weight.shape[0];
        // Compute output spatial dimensions
        let h_out = (input.shape[2] + 2 * padding[0] - dilation[0] * (weight.shape[2] - 1) - 1)
            / stride[0]
            + 1;
        let w_out = (input.shape[3] + 2 * padding[1] - dilation[1] * (weight.shape[3] - 1) - 1)
            / stride[1]
            + 1;
        CubeclTensorPrimitive::new([batch, out_channels, h_out, w_out], input.device)
    }

    // ==================== Pooling Operations ====================

    fn max_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> Self::TensorPrimitive<4> {
        // TODO: Launch max_pool2d kernel
        let batch = input.shape[0];
        let channels = input.shape[1];
        let h_out = (input.shape[2] + 2 * padding[0] - kernel_size[0]) / stride[0] + 1;
        let w_out = (input.shape[3] + 2 * padding[1] - kernel_size[1]) / stride[1] + 1;
        CubeclTensorPrimitive::new([batch, channels, h_out, w_out], input.device)
    }

    fn avg_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
    ) -> Self::TensorPrimitive<4> {
        // TODO: Launch avg_pool2d kernel
        let _ = count_include_pad;
        let batch = input.shape[0];
        let channels = input.shape[1];
        let h_out = (input.shape[2] + 2 * padding[0] - kernel_size[0]) / stride[0] + 1;
        let w_out = (input.shape[3] + 2 * padding[1] - kernel_size[1]) / stride[1] + 1;
        CubeclTensorPrimitive::new([batch, channels, h_out, w_out], input.device)
    }

    // ==================== Device Management ====================

    fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device {
        tensor.device.clone()
    }

    fn to_device<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        // TODO: Copy tensor data to new device if different
        if tensor.device == *device {
            tensor
        } else {
            CubeclTensorPrimitive::new(tensor.shape, device.clone())
        }
    }

    fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
        // Try to read from GPU buffer
        if let Some(buffer) = tensor.buffer() {
            if let Some(gpu_handle) = buffer.as_gpu()
                && let Ok(client) = client::get_or_create(&tensor.device)
                && let Ok(data) = client.read_f32(gpu_handle)
            {
                return data;
            }
            // Read from CPU buffer
            if let Some(bytes) = buffer.as_cpu() {
                return bytes
                    .chunks_exact(4)
                    .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
                    .collect();
            }
        }
        // No buffer - return zeros
        vec![0.0; tensor.numel()]
    }

    fn sync(device: &Self::Device) {
        if let Ok(client) = client::get_or_create(device) {
            client.sync();
        }
    }
}

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

    #[test]
    fn test_backend_name() {
        assert_eq!(CubeclBackend::name(), "CubeCL");
    }

    #[test]
    fn test_zeros() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::zeros::<2>([3, 4], &device);
        assert_eq!(tensor.shape, [3, 4]);
    }

    #[test]
    fn test_matmul_shape() {
        let device = CubeclDevice::cpu();
        let lhs = CubeclBackend::zeros::<2>([3, 4], &device);
        let rhs = CubeclBackend::zeros::<2>([4, 5], &device);
        let result = CubeclBackend::matmul(lhs, rhs);
        assert_eq!(result.shape, [3, 5]);
    }

    #[test]
    fn test_reshape() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::zeros::<2>([3, 4], &device);
        let reshaped = CubeclBackend::reshape::<2, 3>(tensor, [2, 2, 3]);
        assert_eq!(reshaped.shape, [2, 2, 3]);
    }

    #[test]
    fn test_transpose() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::zeros::<2>([3, 4], &device);
        let transposed = CubeclBackend::transpose(tensor);
        assert_eq!(transposed.shape, [4, 3]);
    }

    #[test]
    fn test_device_id() {
        let cpu = CubeclDevice::cpu();
        assert_eq!(cpu.id(), DeviceId::Cpu);

        let cuda = CubeclDevice::new(RuntimeBackend::Cuda, 0);
        assert_eq!(cuda.id(), DeviceId::Cuda(0));
    }

    #[test]
    fn test_from_data_to_data_roundtrip() {
        let device = CubeclDevice::cpu();
        let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
        let tensor = CubeclBackend::from_data(&data, [2, 3], &device);
        let retrieved = CubeclBackend::to_data(&tensor);
        assert_eq!(retrieved.len(), 6);
        // Check values match (using CPU backend)
        for (a, b) in data.iter().zip(retrieved.iter()) {
            assert!((a - b).abs() < 1e-6);
        }
    }

    #[test]
    fn test_zeros_values() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::zeros::<2>([3, 4], &device);
        let data = CubeclBackend::to_data(&tensor);
        assert_eq!(data.len(), 12);
        for val in &data {
            assert!((val - 0.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_ones_values() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::ones::<2>([2, 2], &device);
        let data = CubeclBackend::to_data(&tensor);
        assert_eq!(data.len(), 4);
        for val in &data {
            assert!((val - 1.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_full_values() {
        let device = CubeclDevice::cpu();
        let tensor = CubeclBackend::full::<1>([5], 3.14, &device);
        let data = CubeclBackend::to_data(&tensor);
        assert_eq!(data.len(), 5);
        for val in &data {
            assert!((val - 3.14).abs() < 1e-6);
        }
    }

    #[test]
    fn test_add_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0], [2, 2], &device);
        let b = CubeclBackend::from_data(&[5.0, 6.0, 7.0, 8.0], [2, 2], &device);
        let c = CubeclBackend::add(a, b);
        let data = CubeclBackend::to_data(&c);
        assert_eq!(data, vec![6.0, 8.0, 10.0, 12.0]);
    }

    #[test]
    fn test_mul_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0], [4], &device);
        let b = CubeclBackend::from_data(&[2.0, 3.0, 4.0, 5.0], [4], &device);
        let c = CubeclBackend::mul(a, b);
        let data = CubeclBackend::to_data(&c);
        assert_eq!(data, vec![2.0, 6.0, 12.0, 20.0]);
    }

    #[test]
    fn test_relu_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[-2.0, -1.0, 0.0, 1.0, 2.0], [5], &device);
        let b = CubeclBackend::relu(a);
        let data = CubeclBackend::to_data(&b);
        assert_eq!(data, vec![0.0, 0.0, 0.0, 1.0, 2.0]);
    }

    #[test]
    fn test_sigmoid_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[0.0], [1], &device);
        let b = CubeclBackend::sigmoid(a);
        let data = CubeclBackend::to_data(&b);
        assert!((data[0] - 0.5).abs() < 1e-6);
    }

    #[test]
    fn test_sum_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0], [4], &device);
        let s = CubeclBackend::sum(a);
        let data = CubeclBackend::to_data(&s);
        assert_eq!(data, vec![10.0]);
    }

    #[test]
    fn test_mean_operation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0], [4], &device);
        let m = CubeclBackend::mean(a);
        let data = CubeclBackend::to_data(&m);
        assert!((data[0] - 2.5).abs() < 1e-6);
    }

    #[test]
    fn test_matmul_computation() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3], &device);
        let b = CubeclBackend::from_data(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [3, 2], &device);
        let c = CubeclBackend::matmul(a, b);
        assert_eq!(c.shape, [2, 2]);
        let data = CubeclBackend::to_data(&c);
        // Verify matrix multiplication results
        assert_eq!(data, vec![22.0, 28.0, 49.0, 64.0]);
    }

    #[test]
    fn test_exp_log_roundtrip() {
        let device = CubeclDevice::cpu();
        let a = CubeclBackend::from_data(&[1.0, 2.0, 3.0], [3], &device);
        let b = CubeclBackend::exp(a);
        let c = CubeclBackend::log(b);
        let data = CubeclBackend::to_data(&c);
        for (i, val) in data.iter().enumerate() {
            assert!((val - (i + 1) as f32).abs() < 1e-5);
        }
    }
}