chore(sweep): delete 43 orphaned source files; document SYCL/demo/duplication status

Deletions (all verified unreferenced by any mod/include/path declaration; git history preserves them): - rtx-transformers: entire orphaned curriculum/ split (mod.rs holds the real inline implementation), non-_simple graph variants, superseded simmim/jepa_integration files, layers/{sliding_window_attention, positional_encoding,ssm_state_cache_original}, lib_full/lib_minimal/ error_full/error_minimal, orphaned MoE impls (moe_layer, moe_integration). - rtx-distributed/parallel_old.rs; rtx-flash-attention/{core_full, lib_full}.rs; rtx-compress legacy_distillation + structured_pruner. - rtx-tensor/tensor_core.rs; rtx-runtime/{cuda_kernel_ops, cuda_backend_mock}.rs; rtx-memory/{gpu_pool_manager,allocator, pool_type}.rs; rtx-losses/{lib_minimal,lib_full}.rs. Docs honesty: - rtx-backend-sycl marked EXPERIMENTAL SKELETON in crate docs and CLAUDE.md backend table (all ops return NotImplemented). - docs/consolidation.md records canonical MoE (layers/mixture_of_experts) and flash-attention (rtx-flash-attention crate) implementations plus remaining duplicates to consolidate. - CLAUDE.md: meta-crate GPU features noted; simulation-only demos named; serving/streaming mock removal noted. Verified: cargo check --workspace clean (rtx-onnx-codegen pre-broken at HEAD, unrelated); lib tests pass for all touched crates (rtx-runtime's 4 failures pre-exist at HEAD). Co-Authored-By: Claude Fable 5 <[email protected]>
2026-07-09 19:32:21 -07:00
parent 1e3c604896
commit 5f32165184
47 changed files with 49 additions and 15391 deletions
@@ -106,10 +106,12 @@ let tokens = engine.generate(&prompt, max_tokens=512)?;
 | CUDA | NVIDIA (RTX 4090/5090, A100, H100) | `cuda` |
 | Metal | Apple M-series (M1/M2/M3/M4/M5) | `metal` |
 | ROCm | AMD RDNA 2/3 (RX 7900 XTX) | `rocm` |
-| SYCL | Intel Arc (A770/A750) | `sycl` |
+| SYCL | Intel Arc (A770/A750) — **experimental skeleton, ops return NotImplemented** | `sycl` |
 | WebGPU | Browser (WASM) | `webgpu` |
 | CPU | x86/ARM with MKL | default |
 The meta-crates (`rtx-core`, `rtx-training`, `rtx-inference-stack`) expose `cuda`/`metal` features that thread GPU support through their sub-crates.
 ## Novel Features
 - **Speculative decoding** (`rtx-inference::speculative` + `rtx-inference::medusa`/`lookahead`): 2–3× inference speedup via draft model + target verification — GPU-native Rust implementation
@@ -149,3 +151,5 @@ All in `crates/training/rtx-transformers/src/ssl/`. 163 tests passing.
 ## Current State
 Production-ready. 113 crates. Full transformer stack, Flash Attention (v2+v3), MoE, speculative decoding (Medusa/EAGLE/Lookahead), federated learning, NAS, medical/neuroimaging/scientific domain stacks, complete JEPA platform (Batches 20–26). 30+ demo applications. 13,000+ tests.
 Known honesty notes (2026-07-09 sweep): `rtx-backend-sycl` is an experimental skeleton (ops return NotImplemented). Some demos are pure simulations that don't exercise framework crates (`rtx-distllm-demo`, `rtx-model-zoo`'s MockInferenceEngine, `rtx-inference-profiler`). MoE and flash-attention have duplicated implementations pending consolidation (see `docs/consolidation.md`). Serving/streaming return 503/error until an engine+model is attached (no more mock responses); tokenization at the serving boundary is byte-level until a real tokenizer is threaded through.
@@ -1,6 +1,12 @@
 //! Intel SYCL/oneAPI Backend for RustyTorch++
 //!
-//! This crate provides Intel GPU acceleration using SYCL/oneAPI for:
+//! **STATUS: EXPERIMENTAL SKELETON.** Device queries, buffer create/copy,
 //! and oneMKL GEMM currently return `SyclError::NotImplemented` — this
 //! crate is the port target for a future SYCL implementation and does not
 //! execute on Intel GPUs yet. Use the CUDA, Metal, ROCm, WebGPU, or CPU
 //! backends for real workloads.
 //!
 //! This crate is intended to provide Intel GPU acceleration using SYCL/oneAPI for:
 //! - Intel Arc/Xe discrete GPUs
 //! - Intel integrated GPUs (Gen11+)
 //! - Intel Data Center GPUs (Ponte Vecchio, etc.)
@@ -1,69 +0,0 @@
 //! # RTX Losses - Comprehensive Loss Functions for Deep Learning
 //!
 //! This crate provides a complete collection of loss functions for deep learning applications,
 //! with automatic differentiation support through RTX tensor operations.
 //!
 //! ## Features
 //!
 //! - **FocalLoss**: For imbalanced classification with configurable alpha/gamma parameters
 //! - **DiceLoss**: For segmentation tasks with overlap-based metrics and smoothing
 //! - **ContrastiveLoss**: For metric learning with configurable margin
 //! - **CrossEntropyLoss**: Standard classification with label smoothing support
 //! - **MSELoss**: Regression baseline with robust batch processing
 //! - **InfoNCELoss**: For self-supervised contrastive learning (SimCLR, MoCo)
 //!
 //! All losses support:
 //! - Multiple reduction modes (Mean, Sum, None)
 //! - Automatic gradient computation through RTX autograd
 //! - Numerical stability optimizations
 //! - Comprehensive batch processing
 //!
 //! ## Usage
 //!
 //! ```rust
 //! use rtx_losses::{CrossEntropyLoss, Loss, Reduction};
 //! use rtx_tensor::{Tensor, Device};
 //!
 //! let device = Device::cpu();
 //! let loss = CrossEntropyLoss::new()
 //!     .with_reduction(Reduction::Mean)
 //!     .with_label_smoothing(0.1);
 //!
 //! let logits = Tensor::randn([32, 10], &device)?;  // Batch=32, Classes=10
 //! let targets = Tensor::randint(0, 10, [32], &device)?;
 //!
 //! let loss_value = loss.forward(&logits, &targets)?;
 //! loss_value.backward()?; // Automatic gradient computation
 //! # Ok::<(), rtx_tensor::TensorError>(())
 //! ```
 pub mod error;
 pub mod reduction;
 pub mod loss_trait;
 pub mod focal_loss;
 pub mod dice_loss;
 pub mod contrastive_loss;
 pub mod cross_entropy_loss;
 pub mod mse_loss;
 pub mod info_nce_loss;
 pub mod info_nce_simple;
 pub mod minimal_tensor;
 pub mod info_nce_minimal;
 #[cfg(test)]
 mod standalone_test;
 #[cfg(test)]
 mod simple_test;
 // Re-export main types
 pub use error::{LossError, Result};
 pub use reduction::Reduction;
 pub use loss_trait::Loss;
 pub use focal_loss::FocalLoss;
 pub use dice_loss::DiceLoss;
 pub use contrastive_loss::ContrastiveLoss;
 pub use cross_entropy_loss::CrossEntropyLoss;
 pub use mse_loss::MSELoss;
 pub use info_nce_loss::InfoNCELoss;
 pub use info_nce_simple::InfoNCELossSimple;
 pub use minimal_tensor::{MinimalTensor, MinimalDevice};
 pub use info_nce_minimal::InfoNCEMinimal;
@@ -1,21 +0,0 @@
 //! Minimal RTX Losses - InfoNCE, SwAV, and SimCLR for Self-Supervised Learning
 //! 
 //! This provides working implementations of critical contrastive learning loss functions
 //! with a minimal tensor implementation to avoid ecosystem compilation issues.
 pub mod error;
 pub mod reduction;
 pub mod minimal_tensor;
 pub mod info_nce_minimal;
 // Re-export main types
 pub use error::{LossError, Result};
 pub use reduction::Reduction;
 pub use minimal_tensor::{MinimalTensor, MinimalDevice};
 pub use info_nce_minimal::InfoNCEMinimal;
 // Simplified Loss trait for minimal implementation
 pub trait MinimalLoss {
    /// Reduction mode for this loss
    fn reduction(&self) -> Reduction;
 }
@@ -1,487 +0,0 @@
 //! Memory Allocator Implementation
 //!
 //! Complete memory allocator for GPU and host memory management.
 //! No placeholders, full implementation following strict TDD.
 use crate::error::{MemoryError, Result};
 use crate::pool_type::PoolType;
 use std::collections::{BTreeMap, HashMap};
 use parking_lot::RwLock;
 use tracing::{debug, warn};
 /// Memory allocator for managing memory blocks
 pub struct MemoryAllocator {
    /// Total size of memory managed
    total_size: usize,
    /// Memory alignment requirement
    alignment: usize,
    /// Free blocks tracked by size
    free_blocks: BTreeMap<usize, Vec<MemoryBlock>>,
    /// Allocated blocks tracked by address
    allocated_blocks: HashMap<usize, MemoryBlock>,
    /// Next block ID
    next_block_id: u64,
 }
 /// Memory block descriptor
 #[derive(Debug, Clone, Copy)]
 pub struct MemoryBlock {
    pub id: u64,
    pub address: usize,
    pub size: usize,
    pub allocated: bool,
 }
 /// Allocation information
 #[derive(Debug, Clone)]
 pub struct AllocationInfo {
    pub ptr: usize,
    pub size: usize,
    pub pool_type: PoolType,
    pub timestamp: std::time::Instant,
 }
 impl MemoryAllocator {
    /// Create a new memory allocator
    pub fn new(total_size: usize, alignment: usize) -> Result<Self> {
        if alignment == 0 || !alignment.is_power_of_two() {
            return Err(MemoryError::InvalidConfiguration(
                format!("Alignment must be a power of 2, got {}", alignment)
            ));
        }
        let mut allocator = Self {
            total_size,
            alignment,
            free_blocks: BTreeMap::new(),
            allocated_blocks: HashMap::new(),
            next_block_id: 1,
        };
        // Initialize with one large free block
        let initial_block = MemoryBlock {
            id: 0,
            address: 0,
            size: total_size,
            allocated: false,
        };
        allocator.add_free_block(initial_block);
        debug!("Created allocator with {} bytes, alignment {}", total_size, alignment);
        Ok(allocator)
    }
    /// Allocate memory of the specified size
    pub fn allocate(&mut self, size: usize) -> Result<usize> {
        let aligned_size = self.align_size(size);
        // Find the smallest free block that fits
        let block = self.find_free_block(aligned_size)?;
        // Split the block if it's larger than needed
        if block.size > aligned_size {
            self.split_block(block, aligned_size)?;
        }
        // Mark block as allocated
        let mut allocated_block = block;
        allocated_block.allocated = true;
        allocated_block.size = aligned_size;
        self.allocated_blocks.insert(block.address, allocated_block);
        debug!("Allocated {} bytes at address {:#x}", aligned_size, block.address);
        Ok(block.address)
    }
    /// Deallocate memory at the specified address
    pub fn deallocate(&mut self, address: usize) -> Result<()> {
        let block = self.allocated_blocks.remove(&address)
            .ok_or_else(|| MemoryError::InvalidPointer(
                format!("Address {:#x} not found in allocated blocks", address)
            ))?;
        // Mark as free and try to coalesce with adjacent blocks
        let mut free_block = block;
        free_block.allocated = false;
        self.coalesce_and_add_free(free_block);
        debug!("Deallocated {} bytes at address {:#x}", block.size, address);
        Ok(())
    }
    /// Get total allocated size
    pub fn allocated_size(&self) -> usize {
        self.allocated_blocks.values().map(|b| b.size).sum()
    }
    /// Get total free size
    pub fn free_size(&self) -> usize {
        self.free_blocks
            .values()
            .flat_map(|blocks| blocks.iter())
            .map(|b| b.size)
            .sum()
    }
    /// Get fragmentation ratio
    pub fn fragmentation_ratio(&self) -> f64 {
        let free_block_count = self.free_blocks.values().map(|v| v.len()).sum::<usize>();
        if free_block_count <= 1 {
            0.0
        } else {
            let largest_free = self.free_blocks
                .values()
                .flat_map(|blocks| blocks.iter())
                .map(|b| b.size)
                .max()
                .unwrap_or(0);
            let total_free = self.free_size();
            if total_free == 0 {
                0.0
            } else {
                1.0 - (largest_free as f64 / total_free as f64)
            }
        }
    }
    // Private helper methods
    fn align_size(&self, size: usize) -> usize {
        (size + self.alignment - 1) & !(self.alignment - 1)
    }
    fn find_free_block(&mut self, size: usize) -> Result<MemoryBlock> {
        // Find the smallest block that fits (best fit)
        for (&block_size, blocks) in &mut self.free_blocks {
            if block_size >= size && !blocks.is_empty() {
                return Ok(blocks.remove(0));
            }
        }
        Err(MemoryError::OutOfMemory(format!(
            "No free block found for {} bytes", size
        )))
    }
    fn add_free_block(&mut self, block: MemoryBlock) {
        self.free_blocks
            .entry(block.size)
            .or_default()
            .push(block);
    }
    fn split_block(&mut self, block: MemoryBlock, size: usize) -> Result<()> {
        if block.size <= size {
            return Ok(());
        }
        let remaining_size = block.size - size;
        let remaining_block = MemoryBlock {
            id: self.next_block_id,
            address: block.address + size,
            size: remaining_size,
            allocated: false,
        };
        self.next_block_id += 1;
        self.add_free_block(remaining_block);
        Ok(())
    }
    fn coalesce_and_add_free(&mut self, block: MemoryBlock) {
        let mut coalesced = block;
        // Try to coalesce with previous block
        if coalesced.address > 0 {
            // Check if there's a free block ending at our start address
            for blocks in self.free_blocks.values_mut() {
                if let Some(idx) = blocks.iter().position(|b|
                    b.address + b.size == coalesced.address
                ) {
                    let prev_block = blocks.remove(idx);
                    coalesced.address = prev_block.address;
                    coalesced.size += prev_block.size;
                    break;
                }
            }
        }
        // Try to coalesce with next block
        for blocks in self.free_blocks.values_mut() {
            if let Some(idx) = blocks.iter().position(|b|
                b.address == coalesced.address + coalesced.size
            ) {
                let next_block = blocks.remove(idx);
                coalesced.size += next_block.size;
                break;
            }
        }
        // Clean up empty vectors in free_blocks
        self.free_blocks.retain(|_, blocks| !blocks.is_empty());
        // Add the coalesced block
        self.add_free_block(coalesced);
    }
 }
 /// Buddy allocator for power-of-2 allocations
 pub struct BuddyAllocator {
    min_order: usize,
    max_order: usize,
    free_lists: Vec<Vec<usize>>,
    block_status: HashMap<usize, BlockStatus>,
 }
 #[derive(Debug, Clone, Copy)]
 struct BlockStatus {
    order: usize,
    allocated: bool,
 }
 impl BuddyAllocator {
    /// Create a new buddy allocator
    pub fn new(total_size: usize) -> Result<Self> {
        let max_order = (total_size as f64).log2().floor() as usize;
        let min_order = 12; // 4KB minimum
        if max_order < min_order {
            return Err(MemoryError::InvalidConfiguration(
                format!("Total size too small for buddy allocator")
            ));
        }
        let mut allocator = Self {
            min_order,
            max_order,
            free_lists: vec![Vec::new(); max_order - min_order + 1],
            block_status: HashMap::new(),
        };
        // Add the initial block to the largest order
        allocator.free_lists[max_order - min_order].push(0);
        allocator.block_status.insert(0, BlockStatus {
            order: max_order,
            allocated: false,
        });
        Ok(allocator)
    }
    /// Allocate memory of at least the specified size
    pub fn allocate(&mut self, size: usize) -> Result<usize> {
        let order = self.size_to_order(size);
        if order > self.max_order {
            return Err(MemoryError::OutOfMemory(
                format!("Requested size {} exceeds maximum", size)
            ));
        }
        let address = self.allocate_order(order)?;
        self.block_status.insert(address, BlockStatus {
            order,
            allocated: true,
        });
        Ok(address)
    }
    /// Deallocate memory at the specified address
    pub fn deallocate(&mut self, address: usize) -> Result<()> {
        let status = self.block_status.get(&address)
            .ok_or_else(|| MemoryError::InvalidPointer(
                format!("Address {:#x} not found", address)
            ))?;
        if !status.allocated {
            return Err(MemoryError::InvalidPointer(
                format!("Address {:#x} not allocated", address)
            ));
        }
        let order = status.order;
        self.deallocate_order(address, order);
        Ok(())
    }
    fn size_to_order(&self, size: usize) -> usize {
        let order = (size as f64).log2().ceil() as usize;
        order.max(self.min_order)
    }
    fn allocate_order(&mut self, order: usize) -> Result<usize> {
        // Find a free block of the requested order or split a larger one
        for current_order in order..=self.max_order {
            let idx = current_order - self.min_order;
            if !self.free_lists[idx].is_empty() {
                let address = self.free_lists[idx].remove(0);
                // Split blocks if necessary
                for split_order in (order..current_order).rev() {
                    let buddy_address = address ^ (1 << split_order);
                    self.free_lists[split_order - self.min_order].push(buddy_address);
                    self.block_status.insert(buddy_address, BlockStatus {
                        order: split_order,
                        allocated: false,
                    });
                }
                return Ok(address);
            }
        }
        Err(MemoryError::OutOfMemory(
            format!("No free block for order {}", order)
        ))
    }
    fn deallocate_order(&mut self, address: usize, order: usize) {
        // Try to coalesce with buddy
        let mut current_address = address;
        let mut current_order = order;
        while current_order < self.max_order {
            let buddy_address = current_address ^ (1 << current_order);
            // Check if buddy is free
            if let Some(status) = self.block_status.get(&buddy_address) {
                if !status.allocated && status.order == current_order {
                    // Remove buddy from free list
                    let idx = current_order - self.min_order;
                    if let Some(pos) = self.free_lists[idx].iter().position(|&a| a == buddy_address) {
                        self.free_lists[idx].remove(pos);
                    }
                    // Remove buddy from status
                    self.block_status.remove(&buddy_address);
                    // Continue coalescing at higher order
                    current_address = current_address.min(buddy_address);
                    current_order += 1;
                } else {
                    break;
                }
            } else {
                break;
            }
        }
        // Add the coalesced block to free list
        self.free_lists[current_order - self.min_order].push(current_address);
        self.block_status.insert(current_address, BlockStatus {
            order: current_order,
            allocated: false,
        });
    }
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    #[test]
    fn test_memory_allocator_basic() {
        let mut allocator = MemoryAllocator::new(1024 * 1024, 64).unwrap();
        // Allocate some memory
        let addr1 = allocator.allocate(1024).unwrap();
        assert_eq!(addr1, 0);
        let addr2 = allocator.allocate(2048).unwrap();
        assert_eq!(addr2, 1024);
        // Check sizes
        assert_eq!(allocator.allocated_size(), 1024 + 2048);
        // Deallocate
        allocator.deallocate(addr1).unwrap();
        assert_eq!(allocator.allocated_size(), 2048);
        // Should be able to reuse the space
        let addr3 = allocator.allocate(512).unwrap();
        assert_eq!(addr3, 0);
    }
    #[test]
    fn test_memory_coalescing() {
        let mut allocator = MemoryAllocator::new(1024, 64).unwrap();
        let addr1 = allocator.allocate(256).unwrap();
        let addr2 = allocator.allocate(256).unwrap();
        let addr3 = allocator.allocate(256).unwrap();
        // Deallocate middle block
        allocator.deallocate(addr2).unwrap();
        // Deallocate first block - should not coalesce yet
        allocator.deallocate(addr1).unwrap();
        // Deallocate last block - should coalesce all
        allocator.deallocate(addr3).unwrap();
        // Should be able to allocate the full size again
        let addr4 = allocator.allocate(768).unwrap();
        assert_eq!(addr4, 0);
    }
    #[test]
    fn test_buddy_allocator() {
        let mut buddy = BuddyAllocator::new(1024 * 1024).unwrap();
        // Allocate various sizes
        let addr1 = buddy.allocate(4096).unwrap();
        assert_eq!(addr1, 0);
        let addr2 = buddy.allocate(8192).unwrap();
        assert!(addr2 > 0);
        // Deallocate and check coalescing
        buddy.deallocate(addr1).unwrap();
        buddy.deallocate(addr2).unwrap();
        // Should be able to allocate larger size after coalescing
        let addr3 = buddy.allocate(16384).unwrap();
        assert_eq!(addr3, 0);
    }
    #[test]
    fn test_alignment() {
        let allocator = MemoryAllocator::new(1024, 256).unwrap();
        assert_eq!(allocator.align_size(1), 256);
        assert_eq!(allocator.align_size(256), 256);
        assert_eq!(allocator.align_size(257), 512);
    }
    #[test]
    fn test_fragmentation_ratio() {
        let mut allocator = MemoryAllocator::new(1024, 64).unwrap();
        // Initially no fragmentation
        assert_eq!(allocator.fragmentation_ratio(), 0.0);
        // Create fragmentation
        let addr1 = allocator.allocate(256).unwrap();
        let addr2 = allocator.allocate(256).unwrap();
        let addr3 = allocator.allocate(256).unwrap();
        allocator.deallocate(addr1).unwrap();
        allocator.deallocate(addr3).unwrap();
        // Now we have fragmentation
        let ratio = allocator.fragmentation_ratio();
        assert!(ratio > 0.0);
        assert!(ratio < 1.0);
    }
 }
@@ -1,869 +0,0 @@
 //! GPU Memory Pool Manager
 //!
 //! Integrated GPU memory management combining arena allocation, cross-device transfers,
 //! memory pinning, prefetching, and OOM recovery for production-grade ML workloads.
 use crate::{MemoryError, Result};
 use crate::gpu_allocator::{DeviceId, GpuMemoryBlock, GpuMemoryType, GpuArenaAllocator, GpuMemoryStats};
 use crate::gpu_transfer::{GpuTransferManager, TransferConfig, TransferStats, TransferOperation};
 use crate::gpu_pinning::{GpuPinningManager, PinnedMemoryBlock, PrefetchConfig, PinnedMemoryStats, MemoryAccessType};
 use crate::gpu_oom::{GpuOomManager, OomRecoveryConfig, OomRecoveryStats, AllocatorCallback, MemoryPressure};
 use std::sync::{Arc, Weak};
 use std::collections::HashMap;
 use parking_lot::{RwLock, Mutex};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
 /// Comprehensive GPU memory pool configuration
 #[derive(Debug, Clone)]
 pub struct GpuPoolConfig {
    /// Arena size per device (bytes)
    pub arena_size_per_device: usize,
    /// Memory alignment (must be power of 2)
    pub memory_alignment: usize,
    /// Transfer configuration
    pub transfer_config: TransferConfig,
    /// Prefetch configuration
    pub prefetch_config: PrefetchConfig,
    /// OOM recovery configuration
    pub oom_config: OomRecoveryConfig,
    /// Enable cross-device peer-to-peer transfers
    pub enable_p2p: bool,
    /// Maximum number of devices to manage
    pub max_devices: usize,
    /// Enable performance monitoring
    pub enable_monitoring: bool,
    /// Warmup allocation sizes
    pub warmup_sizes: Vec<usize>,
 }
 impl Default for GpuPoolConfig {
    fn default() -> Self {
        Self {
            arena_size_per_device: 4 * 1024 * 1024 * 1024, // 4GB per device
            memory_alignment: 256, // 256-byte alignment for GPU efficiency
            transfer_config: TransferConfig::default(),
            prefetch_config: PrefetchConfig::default(),
            oom_config: OomRecoveryConfig::default(),
            enable_p2p: true,
            max_devices: 8,
            enable_monitoring: true,
            warmup_sizes: vec![
                4096,      // 4KB - small tensors
                65536,     // 64KB - medium tensors
                1048576,   // 1MB - large tensors
                16777216,  // 16MB - very large tensors
                134217728, // 128MB - model weights
            ],
        }
    }
 }
 /// GPU memory operation types for monitoring
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum GpuMemoryOperation {
    Allocate,
    Deallocate,
    Transfer,
    Pin,
    Unpin,
    Prefetch,
    OomRecover,
 }
 /// Comprehensive GPU memory pool statistics
 #[derive(Debug, Clone)]
 pub struct GpuPoolStats {
    /// Per-device memory statistics
    pub device_stats: HashMap<DeviceId, GpuMemoryStats>,
    /// Transfer statistics
    pub transfer_stats: TransferStats,
    /// Pinned memory statistics
    pub pinned_stats: PinnedMemoryStats,
    /// OOM recovery statistics
    pub oom_stats: OomRecoveryStats,
    /// Total allocations across all devices
    pub total_allocations: u64,
    /// Total memory usage across all devices
    pub total_memory_usage: u64,
    /// Average allocation time in nanoseconds
    pub average_allocation_time_ns: u64,
    /// Cache hit rate across all operations
    pub overall_cache_hit_rate: f64,
    /// Current memory pressure level (worst across devices)
    pub worst_memory_pressure: MemoryPressure,
 }
 /// Unified GPU memory pool manager
 pub struct GpuPoolManager {
    /// Configuration
    config: GpuPoolConfig,
    /// Per-device allocators
    device_allocators: RwLock<HashMap<DeviceId, Arc<GpuArenaAllocator>>>,
    /// Transfer manager for cross-device operations
    transfer_manager: Arc<Mutex<GpuTransferManager>>,
    /// Pinning manager for host-pinned memory
    pinning_manager: Arc<GpuPinningManager>,
    /// OOM recovery manager
    oom_manager: Arc<GpuOomManager>,
    /// Operation statistics
    operation_stats: RwLock<HashMap<GpuMemoryOperation, OperationMetrics>>,
    /// Global allocation counter
    total_allocations: AtomicU64,
    /// Performance monitoring
    monitoring_enabled: bool,
 }
 /// Operation-specific metrics
 #[derive(Debug, Clone)]
 struct OperationMetrics {
    count: u64,
    total_time: Duration,
    success_count: u64,
    bytes_processed: u64,
 }
 impl Default for OperationMetrics {
    fn default() -> Self {
        Self {
            count: 0,
            total_time: Duration::ZERO,
            success_count: 0,
            bytes_processed: 0,
        }
    }
 }
 impl GpuPoolManager {
    /// Create new GPU memory pool manager
    pub fn new(config: GpuPoolConfig) -> Result<Self> {
        let transfer_manager = Arc::new(Mutex::new(
            GpuTransferManager::new(config.transfer_config.clone())
        ));
        let pinning_manager = Arc::new(
            GpuPinningManager::new(config.prefetch_config.clone())
        );
        let oom_manager = Arc::new(
            GpuOomManager::new(config.oom_config.clone())
        );
        let manager = Self {
            config,
            device_allocators: RwLock::new(HashMap::new()),
            transfer_manager,
            pinning_manager,
            oom_manager,
            operation_stats: RwLock::new(HashMap::new()),
            total_allocations: AtomicU64::new(0),
            monitoring_enabled: true,
        };
        Ok(manager)
    }
    /// Initialize GPU device for memory management
    pub fn initialize_device(&self, device_id: DeviceId) -> Result<()> {
        // Check device limit
        if self.device_allocators.read().len() >= self.config.max_devices {
            return Err(MemoryError::configuration(
                format!("Maximum devices limit ({}) reached", self.config.max_devices)
            ));
        }
        // Create device allocator
        let allocator = Arc::new(GpuArenaAllocator::with_alignment(
            device_id,
            self.config.arena_size_per_device,
            self.config.memory_alignment,
        )?);
        // Register with OOM manager
        self.oom_manager.register_allocator(
            Arc::downgrade(&allocator) as Weak<dyn AllocatorCallback>
        );
        // Initialize transfer manager for this device
        {
            let mut transfer_manager = self.transfer_manager.lock();
            transfer_manager.initialize_device(device_id)?;
        }
        // Initialize pinning manager for this device
        #[cfg(feature = "gpu")]
        {
            let mut pinning_manager = Arc::try_unwrap(Arc::clone(&self.pinning_manager))
                .unwrap_or_else(|arc| (*arc).clone());
            pinning_manager.initialize_device(device_id)?;
        }
        // Store the allocator
        {
            let mut allocators = self.device_allocators.write();
            allocators.insert(device_id, allocator);
        }
        // Perform warmup if configured
        if !self.config.warmup_sizes.is_empty() {
            self.warmup_device(device_id)?;
        }
        Ok(())
    }
    /// Allocate GPU memory with automatic device selection and OOM recovery
    pub async fn allocate(&self, size: usize, memory_type: GpuMemoryType, device_id: DeviceId) -> Result<GpuMemoryBlock> {
        let start_time = Instant::now();
        let result = self.try_allocate(size, memory_type, device_id).await;
        // Handle OOM with recovery
        let final_result = match result {
            Err(MemoryError::OutOfMemory { .. }) => {
                // Attempt OOM recovery
                match self.oom_manager.recover_from_oom(size, device_id).await {
                    Ok(true) => {
                        // Recovery successful, retry allocation
                        self.try_allocate(size, memory_type, device_id).await
                    },
                    Ok(false) => {
                        // Recovery failed
                        Err(MemoryError::out_of_memory(size, 0))
                    },
                    Err(recovery_error) => Err(recovery_error),
                }
            },
            other => other,
        };
        // Record operation metrics
        if self.monitoring_enabled {
            let success = final_result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Allocate, start_time.elapsed(), success, size);
        }
        // Track allocation for OOM management
        if let Ok(ref block) = final_result {
            self.oom_manager.track_allocation(*block, None);
            self.total_allocations.fetch_add(1, Ordering::Relaxed);
        }
        final_result
    }
    async fn try_allocate(&self, size: usize, memory_type: GpuMemoryType, device_id: DeviceId) -> Result<GpuMemoryBlock> {
        let allocators = self.device_allocators.read();
        if let Some(allocator) = allocators.get(&device_id) {
            allocator.allocate(size, memory_type)
        } else {
            Err(MemoryError::configuration(
                format!("Device {} not initialized", device_id.id())
            ))
        }
    }
    /// Deallocate GPU memory with tracking cleanup
    pub async fn deallocate(&self, block: GpuMemoryBlock) -> Result<()> {
        let start_time = Instant::now();
        // Untrack from OOM manager
        self.oom_manager.untrack_allocation(block.id);
        // Deallocate from device allocator
        let allocators = self.device_allocators.read();
        let result = if let Some(allocator) = allocators.get(&block.device_id) {
            allocator.deallocate(block)
        } else {
            Err(MemoryError::configuration(
                format!("Device {} not initialized", block.device_id.id())
            ))
        };
        // Record metrics
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Deallocate, start_time.elapsed(), success, block.size);
        }
        result
    }
    /// Allocate host-pinned memory for fast CPU-GPU transfers
    pub async fn allocate_pinned(&self, size: usize, device_id: Option<DeviceId>) -> Result<PinnedMemoryBlock> {
        let start_time = Instant::now();
        let result = self.pinning_manager.allocate_pinned(size, device_id);
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Pin, start_time.elapsed(), success, size);
        }
        result
    }
    /// Free host-pinned memory
    pub async fn free_pinned(&self, block: PinnedMemoryBlock) -> Result<()> {
        let start_time = Instant::now();
        let result = self.pinning_manager.free_pinned(block);
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Unpin, start_time.elapsed(), success, block.size);
        }
        result
    }
    /// Transfer data from CPU to GPU
    pub async fn transfer_host_to_device(&self, src_data: &[u8], dst_block: &GpuMemoryBlock) -> Result<TransferOperation> {
        let start_time = Instant::now();
        let transfer_manager = self.transfer_manager.lock();
        let result = transfer_manager.transfer_host_to_device(src_data, dst_block).await;
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Transfer, start_time.elapsed(), success, src_data.len());
        }
        result
    }
    /// Transfer data from GPU to CPU
    pub async fn transfer_device_to_host(&self, src_block: &GpuMemoryBlock, dst_data: &mut [u8]) -> Result<TransferOperation> {
        let start_time = Instant::now();
        let transfer_manager = self.transfer_manager.lock();
        let result = transfer_manager.transfer_device_to_host(src_block, dst_data).await;
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Transfer, start_time.elapsed(), success, src_block.size);
        }
        result
    }
    /// Transfer data between GPU devices
    pub async fn transfer_device_to_device(&self, src_block: &GpuMemoryBlock, dst_block: &GpuMemoryBlock) -> Result<TransferOperation> {
        let start_time = Instant::now();
        let transfer_manager = self.transfer_manager.lock();
        let result = transfer_manager.transfer_device_to_device(src_block, dst_block).await;
        if self.monitoring_enabled {
            let success = result.is_ok();
            self.record_operation_metric(GpuMemoryOperation::Transfer, start_time.elapsed(), success, src_block.size);
        }
        result
    }
    /// Record memory access pattern for prefetching optimization
    pub fn record_access(&self, block_id: u64, access_type: MemoryAccessType) {
        // Track access for OOM management
        self.oom_manager.track_access(block_id);
        // Track access for prefetching
        self.pinning_manager.record_access(block_id, access_type);
    }
    /// Execute prefetching operations on a device
    pub async fn execute_prefetch(&self, device_id: DeviceId) -> Result<usize> {
        let start_time = Instant::now();
        let result = self.pinning_manager.execute_prefetch(device_id).await;
        if self.monitoring_enabled {
            let success = result.is_ok();
            let prefetched = result.as_ref().unwrap_or(&0);
            self.record_operation_metric(GpuMemoryOperation::Prefetch, start_time.elapsed(), success, *prefetched);
        }
        result
    }
    /// Enable peer-to-peer access between devices
    pub fn enable_peer_access(&self, src_device: DeviceId, dst_device: DeviceId) -> Result<bool> {
        let mut transfer_manager = self.transfer_manager.lock();
        transfer_manager.enable_peer_access(src_device, dst_device)
    }
    /// Check memory pressure across all devices
    pub async fn check_memory_pressure(&self) -> MemoryPressure {
        let allocators = self.device_allocators.read();
        let mut worst_pressure = MemoryPressure::Low;
        for (&device_id, _) in allocators.iter() {
            let pressure = self.oom_manager.check_memory_pressure(device_id).await;
            if pressure > worst_pressure {
                worst_pressure = pressure;
            }
        }
        worst_pressure
    }
    /// Perform proactive memory management across all devices
    pub async fn proactive_memory_management(&self) -> Result<usize> {
        let allocators = self.device_allocators.read();
        let mut total_freed = 0;
        for (&device_id, _) in allocators.iter() {
            match self.oom_manager.proactive_memory_management(device_id).await {
                Ok(freed) => total_freed += freed,
                Err(_) => continue, // Best effort
            }
        }
        Ok(total_freed)
    }
    /// Get comprehensive GPU memory pool statistics
    pub async fn get_stats(&self) -> GpuPoolStats {
        let allocators = self.device_allocators.read();
        let mut device_stats = HashMap::new();
        // Collect per-device statistics
        for (&device_id, allocator) in allocators.iter() {
            device_stats.insert(device_id, allocator.memory_stats());
        }
        // Get component statistics
        let transfer_stats = {
            let transfer_manager = self.transfer_manager.lock();
            transfer_manager.get_transfer_stats()
        };
        let pinned_stats = self.pinning_manager.get_stats();
        let oom_stats = self.oom_manager.get_stats();
        // Calculate aggregated statistics
        let total_memory_usage: u64 = device_stats.values().map(|s| s.current_usage as u64).sum();
        let total_allocations = self.total_allocations.load(Ordering::Relaxed);
        let (average_allocation_time_ns, overall_cache_hit_rate) = self.calculate_performance_metrics();
        let worst_memory_pressure = self.check_memory_pressure().await;
        GpuPoolStats {
            device_stats,
            transfer_stats,
            pinned_stats,
            oom_stats,
            total_allocations,
            total_memory_usage,
            average_allocation_time_ns,
            overall_cache_hit_rate,
            worst_memory_pressure,
        }
    }
    /// Compact memory across all devices to reduce fragmentation
    pub async fn compact_all_devices(&self) -> Result<HashMap<DeviceId, usize>> {
        let allocators = self.device_allocators.read();
        let mut compaction_results = HashMap::new();
        for (&device_id, allocator) in allocators.iter() {
            match allocator.compact() {
                Ok(stats) => {
                    compaction_results.insert(device_id, stats.bytes_recovered);
                },
                Err(_) => {
                    compaction_results.insert(device_id, 0);
                }
            }
        }
        Ok(compaction_results)
    }
    /// Clear all memory caches across devices and managers
    pub async fn clear_all_caches(&self) -> Result<usize> {
        let allocators = self.device_allocators.read();
        let mut total_freed = 0;
        // Clear device allocator caches
        for (_, allocator) in allocators.iter() {
            total_freed += allocator.clear_cache();
        }
        // Clear pinning manager cache
        total_freed += self.pinning_manager.clear_pinned_cache();
        Ok(total_freed)
    }
    /// Reset all statistics
    pub fn reset_stats(&self) {
        {
            let mut stats = self.operation_stats.write();
            stats.clear();
        }
        self.total_allocations.store(0, Ordering::Relaxed);
        {
            let transfer_manager = self.transfer_manager.lock();
            transfer_manager.reset_stats();
        }
        self.oom_manager.reset_stats();
    }
    /// Get memory usage report for all devices
    pub async fn generate_memory_report(&self) -> String {
        let stats = self.get_stats().await;
        let mut report = String::new();
        report.push_str("=== GPU Memory Pool Report ===\n");
        report.push_str(&format!("Total Allocations: {}\n", stats.total_allocations));
        report.push_str(&format!("Total Memory Usage: {} bytes\n", stats.total_memory_usage));
        report.push_str(&format!("Average Allocation Time: {}ns\n", stats.average_allocation_time_ns));
        report.push_str(&format!("Overall Cache Hit Rate: {:.2}%\n", stats.overall_cache_hit_rate * 100.0));
        report.push_str(&format!("Memory Pressure: {:?}\n", stats.worst_memory_pressure));
        report.push_str("\n");
        // Device-specific reports
        for (device_id, device_stats) in &stats.device_stats {
            report.push_str(&format!("Device {}: {} bytes used ({:.1}% utilization)\n",
                device_id.id(),
                device_stats.current_usage,
                device_stats.current_usage as f64 / device_stats.arena_size as f64 * 100.0
            ));
        }
        // Transfer statistics
        report.push_str("\n--- Transfer Statistics ---\n");
        report.push_str(&format!("Total Transfers: {}\n", stats.transfer_stats.total_transfers));
        report.push_str(&format!("Average Bandwidth: {:.2} GB/s\n", stats.transfer_stats.average_bandwidth_gbps));
        report.push_str(&format!("P2P Transfers: {}\n", stats.transfer_stats.p2p_transfers));
        // OOM statistics
        report.push_str("\n--- OOM Recovery Statistics ---\n");
        report.push_str(&format!("Total OOM Events: {}\n", stats.oom_stats.total_oom_events));
        report.push_str(&format!("Successful Recoveries: {}\n", stats.oom_stats.successful_recoveries));
        report.push_str(&format!("Average Recovery Time: {:.2}ms\n", stats.oom_stats.average_recovery_time_ms));
        report
    }
    // Private helper methods
    fn warmup_device(&self, device_id: DeviceId) -> Result<()> {
        let allocators = self.device_allocators.read();
        if let Some(allocator) = allocators.get(&device_id) {
            // Pre-allocate and deallocate common sizes to populate caches
            let mut warmup_blocks = Vec::new();
            for &size in &self.config.warmup_sizes {
                for _ in 0..5 {
                    match allocator.allocate(size, GpuMemoryType::Device) {
                        Ok(block) => warmup_blocks.push(block),
                        Err(_) => break, // If allocation fails, skip this size
                    }
                }
            }
            // Deallocate all blocks to populate free caches
            for block in warmup_blocks {
                let _ = allocator.deallocate(block); // Best effort
            }
        }
        Ok(())
    }
    fn record_operation_metric(&self, operation: GpuMemoryOperation, duration: Duration, success: bool, bytes: usize) {
        let mut stats = self.operation_stats.write();
        let metric = stats.entry(operation).or_default();
        metric.count += 1;
        metric.total_time += duration;
        metric.bytes_processed += bytes as u64;
        if success {
            metric.success_count += 1;
        }
    }
    fn calculate_performance_metrics(&self) -> (u64, f64) {
        let stats = self.operation_stats.read();
        // Calculate average allocation time
        let average_allocation_time_ns = if let Some(alloc_metric) = stats.get(&GpuMemoryOperation::Allocate) {
            if alloc_metric.count > 0 {
                (alloc_metric.total_time.as_nanos() as u64) / alloc_metric.count
            } else {
                0
            }
        } else {
            0
        };
        // Calculate overall cache hit rate (placeholder - would integrate with actual cache metrics)
        let overall_cache_hit_rate = 0.85; // Placeholder value
        (average_allocation_time_ns, overall_cache_hit_rate)
    }
 }
 // Implement AllocatorCallback for integration with OOM manager
 impl AllocatorCallback for GpuArenaAllocator {
    fn get_memory_stats(&self) -> GpuMemoryStats {
        self.memory_stats()
    }
    fn try_free_memory(&self, target_bytes: usize) -> Result<usize> {
        // Try to clear cache first
        let cleared = self.clear_cache();
        if cleared >= target_bytes {
            Ok(cleared)
        } else {
            // Could implement more aggressive freeing strategies here
            Ok(cleared)
        }
    }
    fn compact_memory(&self) -> Result<usize> {
        match self.compact() {
            Ok(stats) => Ok(stats.bytes_recovered),
            Err(e) => Err(e),
        }
    }
    fn clear_caches(&self) -> Result<usize> {
        Ok(self.clear_cache())
    }
    fn get_evictable_blocks(&self) -> Vec<(GpuMemoryBlock, crate::gpu_oom::BlockMetadata)> {
        // This would return blocks that can be evicted
        // For now, return empty as this requires more complex integration
        Vec::new()
    }
    fn evict_blocks(&self, _block_ids: &[u64]) -> Result<usize> {
        // This would evict specific blocks
        // For now, just clear cache as best effort
        Ok(self.clear_cache())
    }
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    #[tokio::test]
    async fn test_gpu_pool_manager_creation() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let stats = manager.get_stats().await;
        assert_eq!(stats.total_allocations, 0);
        assert_eq!(stats.total_memory_usage, 0);
    }
    #[tokio::test]
    async fn test_device_initialization() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        let result = manager.initialize_device(device_id);
        assert!(result.is_ok());
        // Check device is registered
        let allocators = manager.device_allocators.read();
        assert!(allocators.contains_key(&device_id));
    }
    #[tokio::test]
    async fn test_basic_allocation_deallocation() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        // Allocate memory
        let block = manager.allocate(1024, GpuMemoryType::Device, device_id).await.unwrap();
        assert_eq!(block.size, 1024);
        assert_eq!(block.device_id, device_id);
        // Deallocate memory
        let result = manager.deallocate(block).await;
        assert!(result.is_ok());
    }
    #[tokio::test]
    async fn test_pinned_memory_allocation() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        let pinned_block = manager.allocate_pinned(4096, Some(device_id)).await.unwrap();
        assert_eq!(pinned_block.size, 4096);
        assert_eq!(pinned_block.device_id, Some(device_id));
        let result = manager.free_pinned(pinned_block).await;
        assert!(result.is_ok());
    }
    #[tokio::test]
    async fn test_memory_transfer() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        // Allocate GPU memory
        let gpu_block = manager.allocate(1024, GpuMemoryType::Device, device_id).await.unwrap();
        // Test host to device transfer
        let src_data = vec![0xAB; 1024];
        let transfer_result = manager.transfer_host_to_device(&src_data, &gpu_block).await;
        assert!(transfer_result.is_ok());
        // Test device to host transfer
        let mut dst_data = vec![0u8; 1024];
        let transfer_result = manager.transfer_device_to_host(&gpu_block, &mut dst_data).await;
        assert!(transfer_result.is_ok());
        manager.deallocate(gpu_block).await.unwrap();
    }
    #[tokio::test]
    async fn test_memory_pressure_monitoring() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        let pressure = manager.check_memory_pressure().await;
        // Should be low initially
        assert!(matches!(pressure, MemoryPressure::Low | MemoryPressure::Medium));
    }
    #[tokio::test]
    async fn test_access_pattern_recording() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        // Record access patterns
        manager.record_access(1, MemoryAccessType::Sequential);
        manager.record_access(2, MemoryAccessType::Random);
        manager.record_access(3, MemoryAccessType::Strided(8));
        // Should not error
    }
    #[tokio::test]
    async fn test_prefetch_execution() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        let result = manager.execute_prefetch(device_id).await;
        assert!(result.is_ok());
    }
    #[tokio::test]
    async fn test_cache_clearing() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        // Allocate and deallocate to populate caches
        for _ in 0..5 {
            let block = manager.allocate(1024, GpuMemoryType::Device, device_id).await.unwrap();
            manager.deallocate(block).await.unwrap();
        }
        let result = manager.clear_all_caches().await;
        assert!(result.is_ok());
        assert!(result.unwrap() > 0);
    }
    #[tokio::test]
    async fn test_stats_collection() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        // Perform some operations
        let block = manager.allocate(1024, GpuMemoryType::Device, device_id).await.unwrap();
        manager.deallocate(block).await.unwrap();
        let stats = manager.get_stats().await;
        assert!(stats.total_allocations > 0);
        assert!(!stats.device_stats.is_empty());
    }
    #[tokio::test]
    async fn test_memory_report_generation() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        let report = manager.generate_memory_report().await;
        assert!(report.contains("GPU Memory Pool Report"));
        assert!(report.contains("Total Allocations"));
        assert!(report.contains("Device 0"));
    }
    #[tokio::test]
    async fn test_device_limit() {
        let mut config = GpuPoolConfig::default();
        config.max_devices = 1;
        let manager = GpuPoolManager::new(config).unwrap();
        // First device should succeed
        let device1 = DeviceId::new(0);
        assert!(manager.initialize_device(device1).is_ok());
        // Second device should fail due to limit
        let device2 = DeviceId::new(1);
        let result = manager.initialize_device(device2);
        assert!(result.is_err());
    }
    #[tokio::test]
    async fn test_proactive_memory_management() {
        let config = GpuPoolConfig::default();
        let manager = GpuPoolManager::new(config).unwrap();
        let device_id = DeviceId::new(0);
        manager.initialize_device(device_id).unwrap();
        let result = manager.proactive_memory_management().await;
        assert!(result.is_ok());
    }
 }
@@ -1,481 +0,0 @@
 //! Memory Pool Type Implementation
 //!
 //! Complete implementation of different memory pool types for GPU memory management.
 //! No placeholders, full functionality with strict TDD approach.
 use crate::error::{MemoryError, Result};
 use crate::allocator::{MemoryAllocator, AllocationInfo};
 use std::sync::Arc;
 use parking_lot::RwLock;
 use std::collections::HashMap;
 use tracing::{debug, info, warn};
 #[cfg(feature = "cuda")]
 use cudarc::driver::{CudaContext, CudaSlice};
 /// Memory pool type enumeration
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum PoolType {
    /// Device memory (GPU VRAM)
    Device,
    /// Host memory (System RAM)
    Host,
    /// Unified memory (accessible from both CPU and GPU)
    Unified,
    /// Managed memory (automatically migrated between host and device)
    Managed,
    /// Pinned host memory (page-locked for faster transfers)
    Pinned,
 }
 impl PoolType {
    /// Get human-readable name for the pool type
    pub fn name(&self) -> &'static str {
        match self {
            PoolType::Device => "Device",
            PoolType::Host => "Host",
            PoolType::Unified => "Unified",
            PoolType::Managed => "Managed",
            PoolType::Pinned => "Pinned",
        }
    }
    /// Check if this pool type requires GPU support
    pub fn requires_gpu(&self) -> bool {
        match self {
            PoolType::Host => false,
            _ => true,
        }
    }
    /// Get recommended alignment for this pool type
    pub fn alignment(&self) -> usize {
        match self {
            PoolType::Device => 256,  // GPU prefers 256-byte alignment
            PoolType::Unified => 256,
            PoolType::Managed => 256,
            PoolType::Pinned => 64,   // CPU cache line
            PoolType::Host => 64,
        }
    }
    /// Get maximum allocation size hint
    pub fn max_allocation_size(&self) -> usize {
        match self {
            PoolType::Device => 16 * 1024 * 1024 * 1024,  // 16GB typical GPU
            PoolType::Host => usize::MAX,                  // Limited by system
            PoolType::Unified => 8 * 1024 * 1024 * 1024,   // 8GB typical
            PoolType::Managed => 8 * 1024 * 1024 * 1024,
            PoolType::Pinned => 1 * 1024 * 1024 * 1024,    // 1GB pinned is reasonable
        }
    }
 }
 /// Memory pool for managing allocations of a specific type
 pub struct MemoryPool {
    pool_type: PoolType,
    allocator: Arc<RwLock<MemoryAllocator>>,
    allocations: RwLock<HashMap<usize, AllocationInfo>>,
    total_allocated: RwLock<usize>,
    peak_allocated: RwLock<usize>,
    #[cfg(feature = "cuda")]
    cuda_device: Option<Arc<CudaDevice>>,
 }
 impl MemoryPool {
    /// Create a new memory pool
    pub fn new(pool_type: PoolType, initial_size: usize) -> Result<Self> {
        info!("Creating {} memory pool with initial size {}", pool_type.name(), initial_size);
        #[cfg(feature = "cuda")]
        let cuda_device = if pool_type.requires_gpu() {
            Some(Arc::new(CudaDevice::new(0)
                .map_err(|e| MemoryError::InitializationError(format!("Failed to init CUDA: {}", e)))?))
        } else {
            None
        };
        let allocator = Arc::new(RwLock::new(
            MemoryAllocator::new(initial_size, pool_type.alignment())?
        ));
        Ok(Self {
            pool_type,
            allocator,
            allocations: RwLock::new(HashMap::new()),
            total_allocated: RwLock::new(0),
            peak_allocated: RwLock::new(0),
            #[cfg(feature = "cuda")]
            cuda_device,
        })
    }
    /// Allocate memory from the pool
    pub fn allocate(&self, size: usize) -> Result<*mut u8> {
        let aligned_size = align_size(size, self.pool_type.alignment());
        debug!("Allocating {} bytes from {} pool", aligned_size, self.pool_type.name());
        let ptr = match self.pool_type {
            PoolType::Device => self.allocate_device(aligned_size)?,
            PoolType::Host => self.allocate_host(aligned_size)?,
            PoolType::Unified => self.allocate_unified(aligned_size)?,
            PoolType::Managed => self.allocate_managed(aligned_size)?,
            PoolType::Pinned => self.allocate_pinned(aligned_size)?,
        };
        // Track allocation
        let mut allocations = self.allocations.write();
        allocations.insert(ptr as usize, AllocationInfo {
            ptr: ptr as usize,
            size: aligned_size,
            pool_type: self.pool_type,
            timestamp: std::time::Instant::now(),
        });
        // Update statistics
        let mut total = self.total_allocated.write();
        *total += aligned_size;
        let mut peak = self.peak_allocated.write();
        if *total > *peak {
            *peak = *total;
        }
        Ok(ptr)
    }
    /// Deallocate memory back to the pool
    pub fn deallocate(&self, ptr: *mut u8) -> Result<()> {
        let mut allocations = self.allocations.write();
        if let Some(info) = allocations.remove(&(ptr as usize)) {
            debug!("Deallocating {} bytes from {} pool", info.size, self.pool_type.name());
            match self.pool_type {
                PoolType::Device => self.deallocate_device(ptr, info.size)?,
                PoolType::Host => self.deallocate_host(ptr)?,
                PoolType::Unified => self.deallocate_unified(ptr)?,
                PoolType::Managed => self.deallocate_managed(ptr)?,
                PoolType::Pinned => self.deallocate_pinned(ptr)?,
            }
            let mut total = self.total_allocated.write();
            *total = total.saturating_sub(info.size);
            Ok(())
        } else {
            Err(MemoryError::InvalidPointer(format!("Pointer {:p} not found in pool", ptr)))
        }
    }
    /// Get current allocated size
    pub fn allocated_size(&self) -> usize {
        *self.total_allocated.read()
    }
    /// Get peak allocated size
    pub fn peak_allocated_size(&self) -> usize {
        *self.peak_allocated.read()
    }
    /// Get number of active allocations
    pub fn allocation_count(&self) -> usize {
        self.allocations.read().len()
    }
    // Private allocation methods
    #[cfg(feature = "cuda")]
    fn allocate_device(&self, size: usize) -> Result<*mut u8> {
        let device = self.cuda_device.as_ref()
            .ok_or_else(|| MemoryError::InitializationError("CUDA device not initialized".into()))?;
        let slice = device.alloc_zeros::<u8>(size)
            .map_err(|e| MemoryError::AllocationFailed(format!("Device allocation failed: {}", e)))?;
        // Get raw pointer from CudaSlice
        let ptr = slice.device_ptr() as *mut u8;
        // Store the slice to keep it alive
        // In production, we'd need proper lifetime management
        std::mem::forget(slice);
        Ok(ptr)
    }
    #[cfg(not(feature = "cuda"))]
    fn allocate_device(&self, _size: usize) -> Result<*mut u8> {
        Err(MemoryError::UnsupportedOperation("CUDA not enabled".into()))
    }
    fn allocate_host(&self, size: usize) -> Result<*mut u8> {
        let layout = std::alloc::Layout::from_size_align(size, self.pool_type.alignment())
            .map_err(|e| MemoryError::AllocationFailed(format!("Invalid layout: {}", e)))?;
        // SAFETY: alloc allocates memory with the specified layout.
        // - layout is valid (created from size and alignment above)
        // - The resulting pointer is checked for null before use
        let ptr = unsafe { std::alloc::alloc(layout) };
        if ptr.is_null() {
            Err(MemoryError::AllocationFailed(format!("Failed to allocate {} bytes", size)))
        } else {
            Ok(ptr)
        }
    }
    #[cfg(feature = "cuda")]
    fn allocate_unified(&self, size: usize) -> Result<*mut u8> {
        let device = self.cuda_device.as_ref()
            .ok_or_else(|| MemoryError::InitializationError("CUDA device not initialized".into()))?;
        // Use CUDA unified memory
        let slice = device.alloc_unified_zeros::<u8>(size)
            .map_err(|e| MemoryError::AllocationFailed(format!("Unified allocation failed: {}", e)))?;
        let ptr = slice.as_ptr() as *mut u8;
        std::mem::forget(slice);
        Ok(ptr)
    }
    #[cfg(not(feature = "cuda"))]
    fn allocate_unified(&self, _size: usize) -> Result<*mut u8> {
        Err(MemoryError::UnsupportedOperation("CUDA not enabled".into()))
    }
    fn allocate_managed(&self, size: usize) -> Result<*mut u8> {
        // Managed memory is similar to unified in cudarc
        self.allocate_unified(size)
    }
    #[cfg(feature = "cuda")]
    fn allocate_pinned(&self, size: usize) -> Result<*mut u8> {
        let device = self.cuda_device.as_ref()
            .ok_or_else(|| MemoryError::InitializationError("CUDA device not initialized".into()))?;
        // Allocate pinned host memory
        let slice = device.alloc_pinned_zeros::<u8>(size)
            .map_err(|e| MemoryError::AllocationFailed(format!("Pinned allocation failed: {}", e)))?;
        let ptr = slice.as_ptr() as *mut u8;
        std::mem::forget(slice);
        Ok(ptr)
    }
    #[cfg(not(feature = "cuda"))]
    fn allocate_pinned(&self, size: usize) -> Result<*mut u8> {
        // Fall back to regular host allocation
        self.allocate_host(size)
    }
    // Deallocation methods
    #[cfg(feature = "cuda")]
    fn deallocate_device(&self, _ptr: *mut u8, _size: usize) -> Result<()> {
        // In production, we'd properly manage CudaSlice lifetimes
        // For now, memory is freed when the slice is dropped
        Ok(())
    }
    #[cfg(not(feature = "cuda"))]
    fn deallocate_device(&self, _ptr: *mut u8, _size: usize) -> Result<()> {
        Ok(())
    }
    fn deallocate_host(&self, ptr: *mut u8) -> Result<()> {
        if let Some(info) = self.allocations.read().get(&(ptr as usize)) {
            let layout = std::alloc::Layout::from_size_align(info.size, self.pool_type.alignment())
                .map_err(|e| MemoryError::DeallocationFailed(format!("Invalid layout: {}", e)))?;
            // SAFETY: dealloc frees memory previously allocated with alloc.
            // - ptr was obtained from alloc with the same layout (from allocations map)
            // - The layout matches the original allocation
            unsafe {
                std::alloc::dealloc(ptr, layout);
            }
        }
        Ok(())
    }
    fn deallocate_unified(&self, _ptr: *mut u8) -> Result<()> {
        // Unified memory freed when CudaSlice drops
        Ok(())
    }
    fn deallocate_managed(&self, ptr: *mut u8) -> Result<()> {
        self.deallocate_unified(ptr)
    }
    fn deallocate_pinned(&self, ptr: *mut u8) -> Result<()> {
        #[cfg(feature = "cuda")]
        {
            // Pinned memory freed when CudaSlice drops
            Ok(())
        }
        #[cfg(not(feature = "cuda"))]
        {
            self.deallocate_host(ptr)
        }
    }
 }
 /// Align size to the given alignment
 fn align_size(size: usize, alignment: usize) -> usize {
    (size + alignment - 1) & !(alignment - 1)
 }
 /// Global pool manager for different pool types
 pub struct PoolManager {
    pools: RwLock<HashMap<PoolType, Arc<MemoryPool>>>,
 }
 impl PoolManager {
    /// Create a new pool manager
    pub fn new() -> Self {
        Self {
            pools: RwLock::new(HashMap::new()),
        }
    }
    /// Get or create a pool for the given type
    pub fn get_pool(&self, pool_type: PoolType) -> Result<Arc<MemoryPool>> {
        let mut pools = self.pools.write();
        if let Some(pool) = pools.get(&pool_type) {
            Ok(pool.clone())
        } else {
            let initial_size = match pool_type {
                PoolType::Device => 256 * 1024 * 1024,  // 256MB initial
                PoolType::Host => 64 * 1024 * 1024,     // 64MB initial
                _ => 128 * 1024 * 1024,                 // 128MB for others
            };
            let pool = Arc::new(MemoryPool::new(pool_type, initial_size)?);
            pools.insert(pool_type, pool.clone());
            info!("Created new {} pool with {} MB initial size",
                  pool_type.name(), initial_size / (1024 * 1024));
            Ok(pool)
        }
    }
    /// Get statistics for all pools
    pub fn get_statistics(&self) -> HashMap<PoolType, PoolStatistics> {
        let pools = self.pools.read();
        let mut stats = HashMap::new();
        for (pool_type, pool) in pools.iter() {
            stats.insert(*pool_type, PoolStatistics {
                allocated: pool.allocated_size(),
                peak_allocated: pool.peak_allocated_size(),
                allocation_count: pool.allocation_count(),
            });
        }
        stats
    }
 }
 /// Pool statistics
 #[derive(Debug, Clone)]
 pub struct PoolStatistics {
    pub allocated: usize,
    pub peak_allocated: usize,
    pub allocation_count: usize,
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    #[test]
    fn test_pool_type_properties() {
        assert_eq!(PoolType::Device.name(), "Device");
        assert!(PoolType::Device.requires_gpu());
        assert_eq!(PoolType::Device.alignment(), 256);
        assert_eq!(PoolType::Host.name(), "Host");
        assert!(!PoolType::Host.requires_gpu());
        assert_eq!(PoolType::Host.alignment(), 64);
    }
    #[test]
    fn test_align_size() {
        assert_eq!(align_size(100, 64), 128);
        assert_eq!(align_size(64, 64), 64);
        assert_eq!(align_size(1, 256), 256);
        assert_eq!(align_size(257, 256), 512);
    }
    #[test]
    fn test_host_pool_allocation() {
        let pool = MemoryPool::new(PoolType::Host, 1024 * 1024).unwrap();
        let ptr = pool.allocate(1024).unwrap();
        assert!(!ptr.is_null());
        assert_eq!(pool.allocation_count(), 1);
        assert!(pool.allocated_size() >= 1024);
        pool.deallocate(ptr).unwrap();
        assert_eq!(pool.allocation_count(), 0);
    }
    #[test]
    fn test_pool_manager() {
        let manager = PoolManager::new();
        let pool1 = manager.get_pool(PoolType::Host).unwrap();
        let pool2 = manager.get_pool(PoolType::Host).unwrap();
        // Should return the same pool
        assert!(Arc::ptr_eq(&pool1, &pool2));
        let stats = manager.get_statistics();
        assert!(stats.contains_key(&PoolType::Host));
    }
    #[test]
    fn test_multiple_allocations() {
        let pool = MemoryPool::new(PoolType::Host, 10 * 1024 * 1024).unwrap();
        let mut ptrs = Vec::new();
        // Allocate multiple blocks
        for i in 0..10 {
            let size = (i + 1) * 1024;
            let ptr = pool.allocate(size).unwrap();
            assert!(!ptr.is_null());
            ptrs.push(ptr);
        }
        assert_eq!(pool.allocation_count(), 10);
        // Deallocate all
        for ptr in ptrs {
            pool.deallocate(ptr).unwrap();
        }
        assert_eq!(pool.allocation_count(), 0);
    }
    #[test]
    fn test_peak_allocation_tracking() {
        let pool = MemoryPool::new(PoolType::Host, 10 * 1024 * 1024).unwrap();
        let ptr1 = pool.allocate(1024).unwrap();
        let peak1 = pool.peak_allocated_size();
        let ptr2 = pool.allocate(2048).unwrap();
        let peak2 = pool.peak_allocated_size();
        assert!(peak2 > peak1);
        pool.deallocate(ptr1).unwrap();
        let peak3 = pool.peak_allocated_size();
        assert_eq!(peak2, peak3); // Peak should not decrease
        pool.deallocate(ptr2).unwrap();
    }
 }
@@ -1,451 +0,0 @@
 //! Production CUDA Backend Implementation
 //!
 //! This module implements GPU operations using cudarc 0.17.3 safe APIs, replacing
 //! all unsafe sys calls with production-quality safe abstractions. It provides
 //! comprehensive CUDA context management, memory operations, stream handling,
 //! and kernel execution with proper error handling and resource cleanup.
 //!
 //! # Safety
 //!
 //! This implementation eliminates unsafe code by leveraging cudarc's safe APIs.
 //! All GPU operations are properly encapsulated with RAII patterns and comprehensive
 //! error handling using Result types.
 use crate::device::*;
 use crate::error::{RuntimeError, Result};
 use crate::allocator::DevicePtr;
 use std::collections::{BTreeMap};
 use std::sync::{Arc, Mutex, RwLock, atomic::{AtomicBool, Ordering}};
 use tracing::{debug, info, error, warn};
 use anyhow::{Context, bail};
 /// Global CUDA initialization state
 static CUDA_INITIALIZED: AtomicBool = AtomicBool::new(false);
 /// Production CUDA backend with complete resource management
 #[derive(Debug, Clone)]
 pub struct CudaBackend {
    device_id: DeviceId,
    // In a real implementation, this would hold cudarc device handles
    // For now, we'll use simplified placeholders that demonstrate the API
    is_initialized: Arc<AtomicBool>,
 }
 impl CudaBackend {
    /// Create a new CUDA backend for the given device
    ///
    /// This initializes the CUDA device with proper context management and
    /// creates all necessary cuBLAS and cuRAND handles with RAII cleanup.
    pub fn new(device_id: DeviceId) -> Result<Self> {
        initialize_cuda()?;
        // In production, would initialize actual cudarc device here
        info!("Initialized CUDA backend for device {}", device_id.0);
        Ok(Self {
            device_id,
            is_initialized: Arc::new(AtomicBool::new(true)),
        })
    }
    /// Get the device ID
    pub fn device_id(&self) -> DeviceId {
        self.device_id
    }
    /// Allocate GPU memory using cudarc's safe allocator
    ///
    /// This method uses cudarc's built-in memory allocator which provides
    /// proper RAII cleanup and eliminates manual memory management errors.
    pub fn allocate_memory(&self, size: usize) -> Result<DevicePtr> {
        // In production, would use real cudarc memory allocation
        let device_ptr = unsafe { DevicePtr::from_raw((size as u64) + (self.device_id.0 as u64 * 0x1000000)) };
        debug!("Allocated {} bytes at {:#x} on device {}",
               size, device_ptr.as_raw(), self.device_id.0);
        Ok(device_ptr)
    }
    /// Get available GPU memory in bytes
    pub fn get_available_memory(&self) -> Result<usize> {
        // Mock implementation - in production would query actual GPU memory
        let free = 8 * 1024 * 1024 * 1024; // 8GB mock
        debug!("Device {} has {} bytes free memory", self.device_id.0, free);
        Ok(free)
    }
    /// Get total GPU memory in bytes
    pub fn get_total_memory(&self) -> Result<usize> {
        // Mock implementation - in production would query actual GPU memory
        Ok(8 * 1024 * 1024 * 1024) // 8GB mock
    }
    /// Get device properties and capabilities
    pub fn get_device_info(&self) -> Result<DeviceProperties> {
        get_cuda_device_properties(self.device_id.0 as i32)
    }
    /// Create a new CUDA stream
    pub fn create_stream(&self) -> Result<CudaStreamHandle> {
        debug!("Created new CUDA stream on device {}", self.device_id.0);
        Ok(CudaStreamHandle {
            device_id: self.device_id.0 as i32,
            is_valid: Arc::new(AtomicBool::new(true)),
        })
    }
    /// Create a new CUDA event
    pub fn create_event(&self) -> Result<CudaEventHandle> {
        debug!("Created new CUDA event on device {}", self.device_id.0);
        Ok(CudaEventHandle {
            device_id: self.device_id.0 as i32,
            is_valid: Arc::new(AtomicBool::new(true)),
        })
    }
    /// Synchronize the device
    pub fn synchronize(&self) -> Result<()> {
        debug!("Synchronized device {}", self.device_id.0);
        Ok(())
    }
    /// Get cuBLAS handle for matrix operations (placeholder)
    pub fn cublas_handle(&self) -> Option<Arc<Mutex<u64>>> {
        Some(Arc::new(Mutex::new(0x1000))) // Placeholder handle
    }
    /// Get cuRAND handle for random number generation (placeholder)
    pub fn curand_handle(&self) -> Option<Arc<Mutex<u64>>> {
        Some(Arc::new(Mutex::new(0x2000))) // Placeholder handle
    }
 }
 /// Device manager for handling multiple GPU devices with proper resource management
 #[derive(Debug)]
 pub struct DeviceManager {
    devices: RwLock<BTreeMap<DeviceId, Arc<CudaBackend>>>,
    device_count: usize,
 }
 impl DeviceManager {
    /// Create a new device manager and discover all available CUDA devices
    pub fn new() -> Result<Self> {
        initialize_cuda()?;
        // Mock device discovery - in production would use cudarc device enumeration
        let device_count = 1; // Assume at least one device for testing
        info!("Discovered {} CUDA devices", device_count);
        let mut devices = BTreeMap::new();
        // Initialize mock devices
        for i in 0..device_count {
            let device_id = DeviceId(i as u32);
            match CudaBackend::new(device_id) {
                Ok(backend) => {
                    devices.insert(device_id, Arc::new(backend));
                    debug!("Successfully initialized device {}", i);
                }
                Err(e) => {
                    warn!("Failed to initialize device {}: {}", i, e);
                }
            }
        }
        Ok(Self {
            devices: RwLock::new(devices),
            device_count,
        })
    }
    /// Get the number of available devices
    pub fn device_count(&self) -> usize {
        self.device_count
    }
    /// Get a backend for the specified device
    pub fn get_backend(&self, device_id: DeviceId) -> Result<Arc<CudaBackend>> {
        let devices = self.devices.read().unwrap();
        devices.get(&device_id)
            .cloned()
            .ok_or_else(|| RuntimeError::device_error(
                device_id.0,
                format!("Device {} not available", device_id.0)
            ))
    }
    /// Get all available device backends
    pub fn get_all_backends(&self) -> Vec<(DeviceId, Arc<CudaBackend>)> {
        let devices = self.devices.read().unwrap();
        devices.iter().map(|(&id, backend)| (id, backend.clone())).collect()
    }
 }
 impl Default for DeviceManager {
    fn default() -> Self {
        Self::new().expect("Failed to initialize device manager")
    }
 }
 /// CUDA stream wrapper with proper RAII cleanup
 #[derive(Debug)]
 pub struct CudaStreamHandle {
    device_id: i32,
    is_valid: Arc<AtomicBool>,
 }
 impl CudaStreamHandle {
    /// Get the device ID this stream belongs to
    pub fn device_id(&self) -> i32 {
        self.device_id
    }
    /// Synchronize this stream
    pub fn synchronize(&self) -> Result<()> {
        if !self.is_valid.load(Ordering::Acquire) {
            return Err(RuntimeError::stream_error(
                self.device_id as u32,
                "Stream is no longer valid".to_string()
            ));
        }
        debug!("Synchronized stream on device {}", self.device_id);
        Ok(())
    }
    /// Check if stream is completed
    pub fn query(&self) -> Result<bool> {
        if !self.is_valid.load(Ordering::Acquire) {
            return Ok(false);
        }
        // Mock implementation - in production would query actual stream
        Ok(true)
    }
    /// Make this stream wait for an event
    pub fn wait_event(&self, event: &CudaEventHandle) -> Result<()> {
        debug!("Stream on device {} waiting for event", self.device_id);
        Ok(())
    }
 }
 impl Clone for CudaStreamHandle {
    fn clone(&self) -> Self {
        Self {
            device_id: self.device_id,
            is_valid: self.is_valid.clone(),
        }
    }
 }
 /// CUDA event wrapper with proper RAII cleanup
 #[derive(Debug)]
 pub struct CudaEventHandle {
    device_id: i32,
    is_valid: Arc<AtomicBool>,
 }
 impl CudaEventHandle {
    /// Get the device ID this event belongs to
    pub fn device_id(&self) -> i32 {
        self.device_id
    }
    /// Record this event on a stream
    pub fn record(&self, stream: &CudaStreamHandle) -> Result<()> {
        debug!("Recorded event on stream for device {}", self.device_id);
        Ok(())
    }
    /// Synchronize on this event (wait for completion)
    pub fn synchronize(&self) -> Result<()> {
        debug!("Synchronized event on device {}", self.device_id);
        Ok(())
    }
    /// Check if event has completed
    pub fn query(&self) -> Result<bool> {
        if !self.is_valid.load(Ordering::Acquire) {
            return Ok(false);
        }
        Ok(true)
    }
    /// Calculate elapsed time between two events in milliseconds
    pub fn elapsed_time(&self, end_event: &CudaEventHandle) -> Result<f32> {
        // Mock implementation
        Ok(0.5) // 0.5ms mock execution time
    }
 }
 impl Clone for CudaEventHandle {
    fn clone(&self) -> Self {
        Self {
            device_id: self.device_id,
            is_valid: self.is_valid.clone(),
        }
    }
 }
 /// Initialize CUDA driver with proper error handling
 pub fn initialize_cuda() -> Result<()> {
    if CUDA_INITIALIZED.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst).is_err() {
        // Already initialized
        return Ok(());
    }
    // Mock initialization - in production would use cudarc::driver::init()
    info!("CUDA driver initialized successfully (mock implementation)");
    Ok(())
 }
 /// Discover CUDA devices using cudarc safe APIs
 pub fn discover_cuda_devices(devices: &mut BTreeMap<DeviceId, Device>) -> Result<usize> {
    initialize_cuda()?;
    let device_count = 1; // Mock discovery
    info!("Found {} CUDA devices", device_count);
    for cuda_device_id in 0..device_count {
        let device_id = DeviceId(cuda_device_id as u32);
        let properties = get_cuda_device_properties(cuda_device_id as i32)?;
        let device_name = properties.name.clone();
        let device = Device::new(device_id, properties)?;
        devices.insert(device_id, device);
        debug!("Added CUDA device {}: {}", device_id, device_name);
    }
    Ok(device_count)
 }
 /// Get CUDA device properties using mock data (in production would use cudarc)
 fn get_cuda_device_properties(cuda_device_id: i32) -> Result<DeviceProperties> {
    let name = format!("Mock CUDA Device {}", cuda_device_id);
    let total_memory: u64 = 8 * 1024 * 1024 * 1024; // 8GB
    let memory_bandwidth_gb_s = 900.0;
    Ok(DeviceProperties {
        name,
        backend: BackendType::Cuda,
        compute_capability: (8, 9), // SM 89 for RTX 4090/5090
        total_memory: total_memory as u64,
        memory_bandwidth_gb_s,
        multiprocessor_count: 128,
        max_threads_per_block: 1024,
        shared_memory_per_block: 49152,
        warp_size: 32,
        supports_unified_memory: true,
    })
 }
 /// Launch a CUDA kernel with the specified configuration
 ///
 /// # Safety
 /// This function is unsafe as it directly interfaces with CUDA runtime
 pub unsafe fn cuda_launch_kernel(
    module: *const std::ffi::c_void,
    kernel_name: &str,
    grid_dims: (u32, u32, u32),
    block_dims: (u32, u32, u32),
    shared_mem_bytes: u32,
    stream: &CudaStreamHandle,
    params: &[u8],
 ) -> Result<()> {
    // In production, this would use cudarc::driver::launch_kernel
    // For now, we provide a stub implementation
    debug!("Launching kernel {} with grid {:?} block {:?}",
           kernel_name, grid_dims, block_dims);
    if !stream.is_valid.load(Ordering::Acquire) {
        return Err(RuntimeError::kernel_error("Invalid CUDA stream"));
    }
    // Simulate kernel launch - in production would use cudarc
    Ok(())
 }
 /// Load a CUDA module from PTX data
 ///
 /// # Safety
 /// This function is unsafe as it directly loads CUDA code
 pub unsafe fn cuda_module_load_data(ptx_data: *const std::ffi::c_char) -> Result<*const std::ffi::c_void> {
    // In production, this would use cudarc::driver::Module::load_ptx
    // For now, we provide a stub implementation
    debug!("Loading CUDA module from PTX");
    if ptx_data.is_null() {
        return Err(RuntimeError::kernel_error("PTX data is null"));
    }
    // Return a dummy pointer - in production would return actual module handle
    Ok(1 as *const std::ffi::c_void)
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    use std::collections::BTreeMap;
    #[test]
    fn test_cuda_initialization() {
        let result = initialize_cuda();
        assert!(result.is_ok(), "CUDA initialization should succeed");
        // Second call should also succeed (already initialized)
        let result2 = initialize_cuda();
        assert!(result2.is_ok(), "Second CUDA initialization should succeed");
    }
    #[test]
    fn test_device_discovery() {
        let mut devices = BTreeMap::new();
        let count = discover_cuda_devices(&mut devices);
        assert!(count.is_ok(), "Device discovery should succeed");
        let device_count = count.unwrap();
        assert!(device_count > 0, "Should discover at least one CUDA device");
        assert_eq!(devices.len(), device_count, "Device map should contain all discovered devices");
        // Check that we got real device properties
        for device in devices.values() {
            assert_eq!(device.properties.backend, BackendType::Cuda);
            assert!(!device.properties.name.is_empty());
            assert!(device.properties.total_memory > 0);
        }
    }
    #[test]
    fn test_device_manager_creation() {
        let result = DeviceManager::new();
        assert!(result.is_ok(), "DeviceManager creation should succeed");
        let manager = result.unwrap();
        assert!(manager.device_count() > 0, "Should have at least one device");
    }
    #[test]
    fn test_cuda_backend_creation() {
        let result = CudaBackend::new(DeviceId(0));
        assert!(result.is_ok(), "Should be able to create CUDA backend");
        if let Ok(backend) = result {
            assert_eq!(backend.device_id(), DeviceId(0));
            // Test memory operations
            let memory_result = backend.get_available_memory();
            assert!(memory_result.is_ok(), "Should be able to query memory");
            // Test stream creation
            let stream_result = backend.create_stream();
            assert!(stream_result.is_ok(), "Should be able to create stream");
            // Test event creation
            let event_result = backend.create_event();
            assert!(event_result.is_ok(), "Should be able to create event");
        }
    }
 }
@@ -1,554 +0,0 @@
 //! CUDA Kernel Operations using cudarc v0.17.3
 //!
 //! This module provides comprehensive kernel management and execution
 //! using cudarc's safe APIs. No placeholders, full implementations only.
 use crate::error::{RuntimeError, Result};
 use crate::cuda_backend::CudaStreamHandle;
 use std::collections::HashMap;
 use std::sync::{Arc, RwLock};
 use std::sync::atomic::AtomicU64;
 use tracing::{info};
 #[cfg(feature = "cuda")]
 use cudarc::driver::{CudaContext, CudaModule, CudaFunction, CudaStream, LaunchConfig as CudarcLaunchConfig};
 #[cfg(feature = "cuda")]
 use cudarc::nvrtc::{compile_ptx, CompileOptions};
 /// Global module registry for managing loaded PTX modules
 lazy_static::lazy_static! {
    static ref GLOBAL_MODULES: RwLock<HashMap<u64, Arc<LoadedModule>>> = RwLock::new(HashMap::new());
    static ref MODULE_COUNTER: AtomicU64 = AtomicU64::new(1);
 }
 /// Loaded CUDA module with metadata
 #[derive(Debug)]
 pub struct LoadedModule {
    pub id: u64,
    pub name: String,
    #[cfg(feature = "cuda")]
    pub module: CudaModule,
    pub functions: HashMap<String, FunctionInfo>,
    pub ptx_source: String,
    pub compile_time_ms: u64,
 }
 /// Function information within a module
 #[derive(Debug, Clone)]
 pub struct FunctionInfo {
    pub name: String,
    pub max_threads_per_block: u32,
    pub shared_mem_bytes: u32,
    pub const_mem_bytes: u32,
    pub local_mem_bytes: u32,
    pub num_regs: u32,
 }
 /// Kernel launch configuration
 #[derive(Debug, Clone)]
 pub struct KernelLaunchConfig {
    pub grid_dim: (u32, u32, u32),
    pub block_dim: (u32, u32, u32),
    pub shared_mem_bytes: u32,
 }
 impl Default for KernelLaunchConfig {
    fn default() -> Self {
        Self {
            grid_dim: (1, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        }
    }
 }
 /// Launch a CUDA kernel with full parameter support
 ///
 /// This function provides complete kernel launching using cudarc's safe APIs.
 /// It handles PTX compilation, module loading, and parameter marshalling.
 pub fn cuda_launch_kernel(
    _binary_handle: u64,
    _function_name: &str,
    _grid_size: (u32, u32, u32),
    _block_size: (u32, u32, u32),
    _shared_memory_bytes: u32,
    _stream: &CudaStreamHandle,
    _params: &[u8],
 ) -> Result<()> {
    #[cfg(feature = "cuda")]
    {
        // Retrieve the loaded module from global registry
        let modules = GLOBAL_MODULES.read()
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to acquire module lock: {}", e)))?;
        let loaded_module = modules.get(&binary_handle)
            .ok_or_else(|| RuntimeError::kernel_error(format!("Module handle {} not found", binary_handle)))?;
        // Get the function from the module
        let func_info = loaded_module.functions.get(function_name)
            .ok_or_else(|| RuntimeError::kernel_error(format!("Function {} not found in module", function_name)))?;
        // Validate launch configuration
        validate_launch_config(grid_size, block_size, shared_memory_bytes, func_info)?;
        // Get the actual CUDA function
        let cuda_func = loaded_module.module.get_function(function_name)
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to get function: {}", e)))?;
        // Parse and prepare parameters
        let kernel_params = parse_kernel_params(params)?;
        // Build launch configuration
        let config = CudarcLaunchConfig {
            grid_dim: grid_size,
            block_dim: block_size,
            shared_mem_bytes: shared_memory_bytes,
        };
        // Launch the kernel using cudarc's safe API
        let stream_ptr = stream.as_cuda_stream()?;
        unsafe {
            // Launch with parsed parameters
            launch_with_params(&cuda_func, config, &kernel_params, stream_ptr)?;
        }
        debug!("Launched kernel {} with grid {:?}, block {:?}",
               function_name, grid_size, block_size);
        Ok(())
    }
    #[cfg(not(feature = "cuda"))]
    {
        Err(RuntimeError::backend_error(
            "CUDA",
            Box::new(std::io::Error::new(std::io::ErrorKind::Other, "CUDA support not compiled in"))
        ))
    }
 }
 /// Load a PTX module and return a handle
 pub fn cuda_module_load_data(ptx_data: *const std::ffi::c_char) -> Result<u64> {
    #[cfg(feature = "cuda")]
    {
        // Convert C string to Rust string
        let ptx_str = unsafe {
            std::ffi::CStr::from_ptr(ptx_data)
                .to_str()
                .map_err(|e| RuntimeError::kernel_error(format!("Invalid PTX string: {}", e)))?
        };
        // Get the current CUDA device
        let device = get_current_cuda_device()?;
        // Load the PTX module
        let module_name = format!("module_{}", MODULE_COUNTER.fetch_add(1, Ordering::SeqCst));
        let module = device.load_ptx(ptx_str.into(), &module_name, &[])
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to load PTX: {}", e)))?;
        // Extract function information from the module
        let functions = extract_function_info(&module, ptx_str)?;
        // Create module handle
        let handle = MODULE_COUNTER.fetch_add(1, Ordering::SeqCst);
        // Store in global registry
        let loaded_module = Arc::new(LoadedModule {
            id: handle,
            name: module_name,
            module,
            functions,
            ptx_source: ptx_str.to_string(),
            compile_time_ms: 0, // Already compiled
        });
        GLOBAL_MODULES.write()
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to acquire module lock: {}", e)))?
            .insert(handle, loaded_module);
        info!("Loaded PTX module with handle {}", handle);
        Ok(handle)
    }
    #[cfg(not(feature = "cuda"))]
    {
        Err(RuntimeError::backend_error(
            "CUDA",
            Box::new(std::io::Error::new(std::io::ErrorKind::Other, "CUDA support not compiled in"))
        ))
    }
 }
 /// Compile CUDA source code to PTX and load it
 pub fn compile_and_load_cuda(cuda_source: &str, function_names: &[&str]) -> Result<u64> {
    #[cfg(feature = "cuda")]
    {
        let start = std::time::Instant::now();
        // Compile CUDA to PTX using nvrtc
        let compile_opts = CompileOptions {
            arch: Some("sm_70".to_string()), // Minimum for modern GPUs
            include_paths: vec![],
            definitions: vec![],
            disable_warnings: false,
            disable_optimizations: false,
            max_register_count: None,
            use_fast_math: true,
            extra_options: vec![],
        };
        let ptx = compile_ptx(cuda_source, compile_opts)
            .map_err(|e| RuntimeError::kernel_error(format!("PTX compilation failed: {}", e)))?;
        let compile_time_ms = start.elapsed().as_millis() as u64;
        // Load the compiled PTX
        let device = get_current_cuda_device()?;
        let module_name = format!("compiled_{}", MODULE_COUNTER.fetch_add(1, Ordering::SeqCst));
        let module = device.load_ptx(ptx.clone(), &module_name, function_names)
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to load compiled PTX: {}", e)))?;
        // Extract function information
        let mut functions = HashMap::new();
        for func_name in function_names {
            if let Ok(func) = module.get_function(func_name) {
                let info = get_function_attributes(&func)?;
                functions.insert(func_name.to_string(), info);
            }
        }
        // Create module handle
        let handle = MODULE_COUNTER.fetch_add(1, Ordering::SeqCst);
        // Store in global registry
        let loaded_module = Arc::new(LoadedModule {
            id: handle,
            name: module_name.clone(),
            module,
            functions,
            ptx_source: ptx,
            compile_time_ms,
        });
        GLOBAL_MODULES.write()
            .map_err(|e| RuntimeError::kernel_error(format!("Failed to acquire module lock: {}", e)))?
            .insert(handle, loaded_module);
        info!("Compiled and loaded CUDA module {} in {}ms", module_name, compile_time_ms);
        Ok(handle)
    }
    #[cfg(not(feature = "cuda"))]
    {
        Err(RuntimeError::backend_error(
            "CUDA",
            Box::new(std::io::Error::new(std::io::ErrorKind::Other, "CUDA support not compiled in"))
        ))
    }
 }
 /// Get information about a loaded module
 pub fn get_module_info(handle: u64) -> Result<ModuleInfo> {
    let modules = GLOBAL_MODULES.read()
        .map_err(|e| RuntimeError::kernel_error(format!("Failed to acquire module lock: {}", e)))?;
    let module = modules.get(&handle)
        .ok_or_else(|| RuntimeError::kernel_error(format!("Module {} not found", handle)))?;
    Ok(ModuleInfo {
        id: module.id,
        name: module.name.clone(),
        function_count: module.functions.len(),
        function_names: module.functions.keys().cloned().collect(),
        ptx_size: module.ptx_source.len(),
        compile_time_ms: module.compile_time_ms,
    })
 }
 /// Module information structure
 #[derive(Debug, Clone)]
 pub struct ModuleInfo {
    pub id: u64,
    pub name: String,
    pub function_count: usize,
    pub function_names: Vec<String>,
    pub ptx_size: usize,
    pub compile_time_ms: u64,
 }
 /// Unload a module and free resources
 pub fn unload_module(handle: u64) -> Result<()> {
    let mut modules = GLOBAL_MODULES.write()
        .map_err(|e| RuntimeError::kernel_error(format!("Failed to acquire module lock: {}", e)))?;
    if modules.remove(&handle).is_some() {
        info!("Unloaded module {}", handle);
        Ok(())
    } else {
        Err(RuntimeError::kernel_error(format!("Module {} not found", handle)))
    }
 }
 // CudaStreamHandle is imported directly from cuda_backend module
 // Helper functions
 #[cfg(feature = "cuda")]
 fn get_current_cuda_device() -> Result<Arc<CudaDevice>> {
    // Get device 0 for now, can be extended to support multiple devices
    thread_local! {
        static DEVICE: std::cell::RefCell<Option<Arc<CudaDevice>>> = std::cell::RefCell::new(None);
    }
    DEVICE.with(|d| {
        let mut device = d.borrow_mut();
        if device.is_none() {
            *device = Some(Arc::new(
                CudaDevice::new(0)
                    .map_err(|e| RuntimeError::backend_error(format!("Failed to create CUDA device: {}", e)))?
            ));
        }
        Ok(device.as_ref().unwrap().clone())
    })
 }
 #[cfg(feature = "cuda")]
 fn validate_launch_config(
    grid: (u32, u32, u32),
    block: (u32, u32, u32),
    shared_mem: u32,
    func_info: &FunctionInfo,
 ) -> Result<()> {
    let total_threads = block.0 * block.1 * block.2;
    if total_threads > func_info.max_threads_per_block {
        return Err(RuntimeError::config_error(format!(
            "Block size {} exceeds maximum {} for function",
            total_threads, func_info.max_threads_per_block
        )));
    }
    if shared_mem > 48 * 1024 { // 48KB typical shared memory limit
        warn!("Shared memory {} may exceed device limits", shared_mem);
    }
    Ok(())
 }
 #[cfg(feature = "cuda")]
 fn extract_function_info(module: &CudaModule, ptx_source: &str) -> Result<HashMap<String, FunctionInfo>> {
    let mut functions = HashMap::new();
    // Parse PTX to find function declarations
    for line in ptx_source.lines() {
        if line.contains(".entry") || line.contains(".func") {
            if let Some(name) = parse_function_name(line) {
                // Default attributes, can be queried from module if available
                let info = FunctionInfo {
                    name: name.clone(),
                    max_threads_per_block: 1024,
                    shared_mem_bytes: 0,
                    const_mem_bytes: 0,
                    local_mem_bytes: 0,
                    num_regs: 32,
                };
                functions.insert(name, info);
            }
        }
    }
    Ok(functions)
 }
 #[cfg(feature = "cuda")]
 fn parse_function_name(line: &str) -> Option<String> {
    // Find function name in PTX declaration
    if let Some(pos) = line.find(".entry") {
        let after = &line[pos + 6..].trim();
        if let Some(paren) = after.find('(') {
            return Some(after[..paren].trim().to_string());
        }
    }
    if let Some(pos) = line.find(".func") {
        let after = &line[pos + 5..].trim();
        if let Some(paren) = after.find('(') {
            return Some(after[..paren].trim().to_string());
        }
    }
    None
 }
 #[cfg(feature = "cuda")]
 fn get_function_attributes(func: &CudaFunction) -> Result<FunctionInfo> {
    // Query function attributes using cudarc
    // These would come from actual CUDA driver API queries
    Ok(FunctionInfo {
        name: String::new(), // Would be set by caller
        max_threads_per_block: 1024, // Default, query from device
        shared_mem_bytes: 0,
        const_mem_bytes: 0,
        local_mem_bytes: 0,
        num_regs: 32,
    })
 }
 #[cfg(feature = "cuda")]
 fn parse_kernel_params(params: &[u8]) -> Result<Vec<Box<dyn AsKernelParam>>> {
    // Parse byte array into kernel parameters
    // This is a simplified version - real implementation would handle various types
    let mut parsed = Vec::new();
    let mut offset = 0;
    while offset < params.len() {
        // Read parameter type and size from byte stream
        // For now, assume 8-byte aligned parameters
        if offset + 8 <= params.len() {
            let param_bytes = &params[offset..offset + 8];
            let value = u64::from_le_bytes(param_bytes.try_into().unwrap());
            parsed.push(Box::new(value) as Box<dyn AsKernelParam>);
            offset += 8;
        } else {
            break;
        }
    }
    Ok(parsed)
 }
 #[cfg(feature = "cuda")]
 fn launch_with_params(
    func: &CudaFunction,
    config: CudarcLaunchConfig,
    params: &[Box<dyn AsKernelParam>],
    stream: &CudaStream,
 ) -> Result<()> {
    // Use cudarc's launch builder
    let mut builder = func.launch_builder();
    for param in params {
        builder.arg(param.as_ref());
    }
    unsafe {
        builder.launch_on_stream(config, stream)
            .map_err(|e| RuntimeError::kernel_error(format!("Kernel launch failed: {}", e)))?;
    }
    Ok(())
 }
 // Trait for kernel parameters
 trait AsKernelParam {
    fn as_kernel_param(&self) -> *const std::ffi::c_void;
 }
 impl AsKernelParam for u32 {
    fn as_kernel_param(&self) -> *const std::ffi::c_void {
        self as *const _ as *const std::ffi::c_void
    }
 }
 impl AsKernelParam for u64 {
    fn as_kernel_param(&self) -> *const std::ffi::c_void {
        self as *const _ as *const std::ffi::c_void
    }
 }
 impl AsKernelParam for f32 {
    fn as_kernel_param(&self) -> *const std::ffi::c_void {
        self as *const _ as *const std::ffi::c_void
    }
 }
 impl AsKernelParam for f64 {
    fn as_kernel_param(&self) -> *const std::ffi::c_void {
        self as *const _ as *const std::ffi::c_void
    }
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    #[test]
    fn test_kernel_launch_config_default() {
        let config = KernelLaunchConfig::default();
        assert_eq!(config.grid_dim, (1, 1, 1));
        assert_eq!(config.block_dim, (256, 1, 1));
        assert_eq!(config.shared_mem_bytes, 0);
    }
    #[test]
    fn test_module_info_creation() {
        let info = ModuleInfo {
            id: 1,
            name: "test_module".to_string(),
            function_count: 2,
            function_names: vec!["func1".to_string(), "func2".to_string()],
            ptx_size: 1024,
            compile_time_ms: 100,
        };
        assert_eq!(info.function_count, 2);
        assert_eq!(info.function_names.len(), 2);
    }
    #[cfg(feature = "cuda")]
    #[test]
    fn test_parse_function_name() {
        let line1 = ".entry vector_add(";
        assert_eq!(parse_function_name(line1), Some("vector_add".to_string()));
        let line2 = ".func matrix_mul(";
        assert_eq!(parse_function_name(line2), Some("matrix_mul".to_string()));
        let line3 = "// Just a comment";
        assert_eq!(parse_function_name(line3), None);
    }
    #[test]
    fn test_validate_module_handle() {
        // Test that invalid handles return errors
        let result = get_module_info(99999);
        assert!(result.is_err());
    }
    #[cfg(feature = "cuda")]
    #[test]
    fn test_compile_simple_kernel() {
        let cuda_source = r#"
            extern "C" __global__ void add_one(float* data, int n) {
                int idx = blockIdx.x * blockDim.x + threadIdx.x;
                if (idx < n) {
                    data[idx] += 1.0f;
                }
            }
        "#;
        let result = compile_and_load_cuda(cuda_source, &["add_one"]);
        // This will fail without CUDA, but tests compilation logic
        if cfg!(feature = "cuda") {
            match result {
                Ok(handle) => {
                    // Verify module was loaded
                    let info = get_module_info(handle);
                    assert!(info.is_ok());
                    // Clean up
                    let _ = unload_module(handle);
                },
                Err(e) => {
                    // Expected if no CUDA device available
                    println!("CUDA compilation test skipped: {}", e);
                }
            }
        } else {
            assert!(result.is_err());
        }
    }
 }
@@ -1,815 +0,0 @@
 //! Core Tensor Structure and Fundamental Operations
 //!
 //! This module provides the main `Tensor` type and its core functionality including:
 //! - Tensor structure definition
 //! - Constructor functions (zeros, ones, eye, from_data)
 //! - Basic properties (shape, dtype, device, etc.)
 //! - Memory management and views
 //! - Autograd integration infrastructure
 use crate::{Device, DType, Shape, Storage, TensorError, Result};
 use std::sync::{Arc, Mutex};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use rand::Rng;
 /// Global atomic counter for autograd node IDs
 static NEXT_NODE_ID: AtomicUsize = AtomicUsize::new(0);
 /// Secondary atomic counter for tensor creation with requires_grad
 static NEXT_GRAD_NODE_ID: AtomicUsize = AtomicUsize::new(1000);
 /// Simple node ID for autograd integration (temporary until proper integration)
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct NodeId(pub usize);
 impl NodeId {
    pub fn new(id: usize) -> Self {
        NodeId(id)
    }
    /// Get the underlying ID value
    pub fn id(&self) -> usize {
        self.0
    }
 }
 /// Core tensor type with GPU memory backing and PyTorch-compatible API
 ///
 /// # Design Principles
 /// - GPU memory backed by rtx-runtime allocator
 /// - Reference counting for efficient memory usage
 /// - Zero-copy views via shared storage with offset/stride
 /// - Automatic gradient computation support (Phase 2)
 ///
 /// # Field Ordering (Cache Locality Optimization)
 /// Fields are ordered by access pattern to maximize cache locality:
 /// 1. Hot path (memory access): storage, strides, offset
 /// 2. Hot path (compatibility): shape, device
 /// 3. Warm (autograd check): dtype, requires_grad
 /// 4. Cold (backward only): node_id, grad
 #[derive(Debug, Clone)]
 pub struct Tensor {
    // === Hot: Memory access fields (accessed together for data operations) ===
    /// Underlying storage with GPU memory
    storage: Arc<Storage>,
    /// Memory stride for each dimension
    strides: Vec<usize>,
    /// Offset into storage buffer
    offset: usize,
    // === Hot: Shape/type/device (compatibility checks on every operation) ===
    /// Tensor shape (dimensions)
    shape: Shape,
    /// Device location
    device: Device,
    // === Warm: Type and autograd flag (checked frequently) ===
    /// Data type
    dtype: DType,
    /// Whether gradient computation is required
    requires_grad: bool,
    // === Cold: Backward pass only (rarely accessed during forward) ===
    /// Node ID in computation graph for autograd
    node_id: Option<NodeId>,
    /// Gradient tensor (lazily allocated)
    grad: Arc<Mutex<Option<Tensor>>>,
 }
 impl Tensor {
    /// Create a new tensor with zeros
    ///
    /// # Arguments
    /// * `shape` - Tensor dimensions
    /// * `device` - Target device for memory allocation
    ///
    /// # Returns
    /// New zero-initialized tensor
    pub fn zeros<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
        let shape = shape.into();
        let numel = shape.numel();
        let storage = Arc::new(Storage::zeros(numel, DType::F32, device)?);
        let strides = shape.default_strides();
        Ok(Tensor {
            storage,
            shape,
            strides,
            offset: 0,
            dtype: DType::F32,
            device: device.clone(),
            requires_grad: false,
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Create a new tensor with ones
    pub fn ones<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
        let shape = shape.into();
        let numel = shape.numel();
        let storage = Arc::new(Storage::full(numel, 1.0, DType::F32, device)?);
        let strides = shape.default_strides();
        Ok(Tensor {
            storage,
            shape,
            strides,
            offset: 0,
            dtype: DType::F32,
            device: device.clone(),
            requires_grad: false,
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Create an identity matrix
    pub fn eye(n: usize, device: &Device) -> Result<Self> {
        let mut data = vec![0.0f32; n * n];
        for i in 0..n {
            data[i * n + i] = 1.0;
        }
        let storage = Arc::new(Storage::from_data(data, DType::F32, device)?);
        let shape = Shape::from(vec![n, n]);
        let strides = shape.default_strides();
        Ok(Tensor {
            storage,
            shape,
            strides,
            offset: 0,
            dtype: DType::F32,
            device: device.clone(),
            requires_grad: false,
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Create a tensor from raw data
    pub fn from_data<S: Into<Shape>>(
    data: Vec<f32>,
    shape: S,
    device: &Device,
    ) -> Result<Self> {
        let shape = shape.into();
        if data.len() != shape.numel() {
            return Err(TensorError::shape(format!(
            "Data length {} doesn't match shape size {}",
            data.len(),
            shape.numel()
            )));
        }
        let storage = Arc::new(Storage::from_vec(data, DType::F32, device)?);
        let strides = shape.default_strides();
        Ok(Tensor {
            storage,
            shape,
            strides,
            offset: 0,
            dtype: DType::F32,
            device: device.clone(),
            requires_grad: false,
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Create a tensor with random values from uniform distribution [0, 1)
    pub fn rand<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
        let shape = shape.into();
        let numel = shape.numel();
        let mut rng = rand::thread_rng();
        let data: Vec<f32> = (0..numel).map(|_| rng.r#gen()).collect();
        Self::from_data(data, shape, device)
    }
    /// Create a tensor with random values from standard normal distribution
    pub fn randn<S: Into<Shape>>(shape: S, device: &Device) -> Result<Self> {
        let shape = shape.into();
        let numel = shape.numel();
        let mut rng = rand::thread_rng();
        let data: Vec<f32> = (0..numel)
        .map(|_| {
            // Box-Muller transform for normal distribution
            let u1: f32 = rng.r#gen();
            let u2: f32 = rng.r#gen();
            (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos()
        })
        .collect();
        Self::from_data(data, shape, device)
    }
    /// Get tensor shape
    ///
    /// Marked inline as called frequently in hot paths (shape checks, broadcasting).
    #[inline]
    pub fn shape(&self) -> &Shape {
        &self.shape
    }
    /// Get number of dimensions
    ///
    /// Marked inline as frequently chained with shape access.
    #[inline]
    pub fn ndim(&self) -> usize {
        self.shape.ndim()
    }
    /// Get total number of elements
    ///
    /// Marked inline as used in tight loops for size calculations.
    #[inline]
    pub fn numel(&self) -> usize {
        self.shape.numel()
    }
    /// Get device location
    ///
    /// Marked inline as called for every cross-device check.
    #[inline]
    pub fn device(&self) -> &Device {
        &self.device
    }
    /// Get data type
    ///
    /// Marked inline as called for every type compatibility check.
    #[inline]
    pub fn dtype(&self) -> DType {
        self.dtype
    }
    /// Check if gradients are required
    ///
    /// Marked inline as checked frequently during forward/backward passes.
    #[inline]
    pub fn requires_grad(&self) -> bool {
        self.requires_grad
    }
    /// Set gradient requirement
    pub fn set_requires_grad(&mut self, requires_grad: bool) {
        self.requires_grad = requires_grad;
    }
    /// Get the gradient tensor
    pub fn grad(&self) -> Option<Tensor> {
        // Use unwrap_or_else to recover data even if mutex was poisoned by another thread
        self.grad.lock().unwrap_or_else(|poisoned| poisoned.into_inner()).clone()
    }
    /// Set the gradient tensor
    pub fn set_grad(&self, grad: Option<Tensor>) {
        // Use unwrap_or_else to recover data even if mutex was poisoned by another thread
        *self.grad.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = grad;
    }
    /// Get the node ID for autograd
    pub fn node_id(&self) -> Option<NodeId> {
        self.node_id
    }
    /// Set the node ID for autograd (internal use)
    pub(crate) fn set_node_id(&mut self, node_id: NodeId) {
        self.node_id = Some(node_id);
    }
    /// Create a tensor that requires gradients and has a node ID
    pub fn require_grad(mut self) -> Self {
        self.requires_grad = true;
        // Node ID will be set when added to tape
        self
    }
    /// Get or create a node ID for this tensor (simplified for now)
    pub(crate) fn get_or_create_node_id(&mut self) -> Result<NodeId> {
        if let Some(node_id) = self.node_id {
            return Ok(node_id);
        }
        // Use atomic counter for thread-safe node ID generation
        let node_id = NodeId::new(NEXT_NODE_ID.fetch_add(1, Ordering::Relaxed));
        self.node_id = Some(node_id);
        Ok(node_id)
    }
    /// Compute gradients via backward pass starting from this tensor (placeholder)
    pub fn backward(&self) -> Result<()> {
        if self.node_id.is_none() {
            return Err(TensorError::autograd("Cannot backward: tensor has no node ID".to_string()));
        }
        // This is a placeholder implementation
        // The actual backward pass will be implemented in the autograd integration layer
        // For now, just return success to allow compilation
        Ok(())
    }
    /// Create a view of this tensor with new shape
    pub fn view<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
        let new_shape = shape.into();
        // Check that the new shape has the same number of elements
        if !self.shape.can_reshape_to(&new_shape) {
            return Err(TensorError::shape(format!(
            "Cannot view tensor with {} elements as shape with {} elements",
            self.shape.numel(),
            new_shape.numel()
            )));
        }
        // Check if tensor is contiguous for zero-copy view
        // For simplicity, we assume contiguous layout for now
        // In practice, this would check strides vs. default strides
        let new_strides = new_shape.default_strides();
        // Create view sharing the same storage
        Ok(Tensor {
            storage: self.storage.clone(),
            shape: new_shape,
            strides: new_strides,
            offset: self.offset,
            dtype: self.dtype,
            device: self.device.clone(),
            requires_grad: self.requires_grad,
            grad: Arc::new(Mutex::new(None)), // New tensor starts with no gradient
            node_id: None, // Views get new node IDs when used in operations
        })
    }
    /// Reshape tensor (may require copy if not contiguous)
    pub fn reshape<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
        let new_shape = shape.into();
        // Check that the new shape has the same number of elements
        if !self.shape.can_reshape_to(&new_shape) {
            return Err(TensorError::shape(format!(
            "Cannot reshape tensor with {} elements to shape with {} elements",
            self.shape.numel(),
            new_shape.numel()
            )));
        }
        // For now, always create a copy since we're not tracking contiguity
        // In practice, this would try view() first and fall back to copy if needed
        let new_numel = new_shape.numel();
        let new_storage = Arc::new(Storage::zeros(new_numel, self.dtype, &self.device)?);
        let new_strides = new_shape.default_strides();
        // Copy data from current tensor
        let data = self.to_cpu()?;
        let mut result_tensor = Tensor {
            storage: new_storage,
            shape: new_shape,
            strides: new_strides,
            offset: 0,
            dtype: self.dtype,
            device: self.device.clone(),
            requires_grad: self.requires_grad,
            grad: Arc::new(Mutex::new(None)), // New tensor starts with no gradient
            node_id: None, // Reshaping creates new node ID when used in operations
        };
        Arc::get_mut(&mut result_tensor.storage)
            .expect("Internal error: newly created tensor should have exclusive storage access")
            .copy_from_cpu(&data)?;
        Ok(result_tensor)
    }
    /// Move tensor to CPU and return data as Vec<f32>
    pub fn to_cpu(&self) -> Result<Vec<f32>> {
        self.storage.to_cpu()
    }
    /// Move tensor to specified device
    pub fn to_device(&self, device: &Device) -> Result<Self> {
        if device == &self.device {
            return Ok(self.clone());
        }
        // Transfer data via CPU (simplified implementation)
        let data = self.to_cpu()?;
        let new_storage = Arc::new(Storage::from_vec(data, self.dtype, device)?);
        Ok(Tensor {
            storage: new_storage,
            shape: self.shape.clone(),
            strides: self.strides.clone(),
            offset: self.offset,
            dtype: self.dtype,
            device: device.clone(),
            requires_grad: self.requires_grad,
            grad: Arc::new(Mutex::new(None)), // New tensor starts with no gradient
            node_id: None, // Device transfer creates new node ID when used in operations
        })
    }
    /// Get the CUDA device pointer (`CUdeviceptr`) for use with GPU kernel launches.
    ///
    /// Returns `Err` when the tensor is not on a CUDA device.  The returned
    /// value is a raw GPU virtual address (a `u64`) that can be passed to
    /// `cudarc`'s `launch_builder().arg(&ptr)` as a kernel argument.
    #[cfg(feature = "cuda")]
    pub fn cuda_device_ptr(&self) -> Result<cudarc::driver::sys::CUdeviceptr> {
        self.storage.cuda_device_ptr()
    }
    /// Get raw data pointer for CUDA kernel access
    ///
    /// # Safety
    /// The returned pointer is only valid as long as the tensor exists and is not modified.
    /// This is primarily for interfacing with CUDA kernels.
    pub unsafe fn data_ptr(&self) -> *const u8 {
        // SAFETY: Pointer access for CUDA kernel interfacing:
        // - Caller guarantees tensor lifetime and no modifications per API contract
        // - Storage is cloned to avoid borrow checker conflicts
        // - Underlying GPU memory is valid (managed by Arc<Storage>)
        // - Used primarily for passing pointers to CUDA kernels
        let mut storage_clone = (*self.storage).clone();
        storage_clone.data_ptr()
    }
    /// Get mutable raw data pointer for CUDA kernel access
    ///
    /// # Safety
    /// The returned pointer is only valid as long as the tensor exists.
    /// Modifications through this pointer will affect the tensor data.
    /// This is primarily for interfacing with CUDA kernels.
    pub unsafe fn data_ptr_mut(&mut self) -> *mut u8 {
        // SAFETY: Mutable pointer access for CUDA kernel interfacing:
        // - Caller guarantees exclusive access and proper lifetime per API contract
        // - Copy-on-write: if Arc refcount > 1, storage is cloned first
        // - After clone, Arc::get_mut is guaranteed to succeed (refcount == 1)
        // - Underlying GPU memory is valid (managed by Arc<Storage>)
        // Get mutable access to the storage through Arc
        Arc::get_mut(&mut self.storage)
            .map(|s| s.data_ptr_mut())
            .unwrap_or_else(|| {
                // If we can't get exclusive access, clone the storage
                let new_storage = (*self.storage).clone();
                self.storage = Arc::new(new_storage);
                Arc::get_mut(&mut self.storage)
                    .expect("Internal error: freshly cloned Arc should have exclusive access")
                    .data_ptr_mut()
            })
    }
    /// Convert tensor to specified data type
    pub fn to_dtype(&self, dtype: DType) -> Result<Self> {
        if dtype == self.dtype {
            return Ok(self.clone());
        }
        // For now, only support f32 (simplified implementation)
        match dtype {
            DType::F32 => Ok(self.clone()),
            _ => Err(TensorError::type_error(format!(
            "Data type conversion to {:?} not yet implemented",
            dtype
            ))),
        }
    }
    /// Get a scalar value (for tensors with single element)
    pub fn item(&self) -> Result<f32> {
        if self.numel() != 1 {
            return Err(TensorError::shape(format!(
            "item() can only be called on tensors with a single element, got {} elements",
            self.numel()
            )));
        }
        let data = self.to_cpu()?;
        Ok(data[0])
    }
    /// Clone tensor with new storage (deep copy)
    pub fn clone_detached(&self) -> Result<Self> {
        let data = self.to_cpu()?;
        let new_storage = Arc::new(Storage::from_vec(data, self.dtype, &self.device)?);
        Ok(Tensor {
            storage: new_storage,
            shape: self.shape.clone(),
            strides: self.strides.clone(),
            offset: self.offset,
            dtype: self.dtype,
            device: self.device.clone(),
            requires_grad: false, // Detached tensors don't require gradients
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Utility functions for indexing and broadcasting
    /// Convert flat index to multi-dimensional coordinates
    pub(crate) fn unravel_index(index: usize, shape: &Shape) -> Vec<usize> {
        let dims = shape.dims();
        let mut coords = Vec::with_capacity(dims.len());
        let mut remaining = index;
        for &dim_size in dims.iter().rev() {
            coords.push(remaining % dim_size);
            remaining /= dim_size;
        }
        coords.reverse();
        coords
    }
    /// Convert multi-dimensional coordinates to flat index
    pub(crate) fn ravel_index(coords: &[usize], strides: &[usize]) -> usize {
        coords.iter().zip(strides.iter()).map(|(coord, stride)| coord * stride).sum()
    }
    /// Compute broadcast index for given coordinates
    pub(crate) fn broadcast_index(coords: &[usize], src_shape: &Shape, _result_shape: &Shape) -> usize {
        let src_dims = src_shape.dims();
        let src_strides = src_shape.default_strides();
        // Handle broadcasting by mapping coordinates to source tensor dimensions
        let mut src_coords = Vec::with_capacity(src_dims.len());
        let offset = coords.len() - src_dims.len();
        for (i, &dim_size) in src_dims.iter().enumerate() {
            let coord_idx = offset + i;
            if coord_idx < coords.len() {
                let coord = coords[coord_idx];
                src_coords.push(if dim_size == 1 { 0 } else { coord });
                } else {
                    src_coords.push(0);
                }
            }
            Self::ravel_index(&src_coords, &src_strides)
        }
        /// Internal utility to create a tensor with automatic node ID assignment
        pub(crate) fn create_with_autograd(
        storage: Arc<Storage>,
        shape: Shape,
        dtype: DType,
        device: Device,
        requires_grad: bool,
        base_node_id: usize,
        ) -> Self {
            let mut tensor = Tensor {
                storage,
                shape,
                strides: shape.default_strides(),
                offset: 0,
                dtype,
                device,
                requires_grad,
                grad: Arc::new(Mutex::new(None)),
                node_id: None,
            };
            if requires_grad {
                let id = base_node_id + NEXT_GRAD_NODE_ID.fetch_add(1, Ordering::Relaxed);
                tensor.node_id = Some(NodeId::new(id));
            }
            tensor
        }
        /// Advanced tensor core optimization utilities
        pub fn optimize_for_tensor_cores(&self) -> Result<TensorCoreOptimization> {
            let shape = self.shape().dims();
            let device = self.device();
            if !device.is_gpu() {
                return Ok(TensorCoreOptimization {
                    can_use_tensor_cores: false,
                    recommended_precision: MixedPrecision::FP32,
                    optimal_tile_size: None,
                    memory_layout: MemoryLayout::RowMajor,
                });
            }
            // Check tensor core compatibility
            let can_use_tensor_cores = check_tensor_core_compatibility(shape, self.dtype());
            let recommended_precision = if can_use_tensor_cores {
                select_optimal_precision(shape, self.dtype())
            } else {
                MixedPrecision::FP32
            };
            Ok(TensorCoreOptimization {
                can_use_tensor_cores,
                recommended_precision,
                optimal_tile_size: None, // Simplified for now
                memory_layout: MemoryLayout::RowMajor,
            })
        }
        /// Mixed precision matrix multiplication
        pub fn mixed_precision_matmul(&self, other: &Tensor, precision: MixedPrecision) -> Result<Tensor> {
            // For now, all precision modes fall back to FP32
            // Full implementation would handle actual precision conversion
            self.matmul(other)
        }
    // =========================================================================
    // IN-PLACE OPERATION SUPPORT (Copy-on-Write)
    // =========================================================================
    /// Check if this tensor has exclusive ownership of its storage.
    /// Returns true if safe to mutate in-place (refcount == 1).
    pub fn is_exclusive(&self) -> bool {
        Arc::strong_count(&self.storage) == 1
    }
    /// Ensure exclusive ownership of storage by cloning if needed (Copy-on-Write).
    /// After this call, in-place modifications are safe.
    pub fn make_exclusive(&mut self) -> Result<()> {
        if !self.is_exclusive() {
            // Clone the storage data
            let data = self.to_cpu()?;
            let new_storage = Storage::from_vec(data, self.dtype, &self.device)?;
            self.storage = Arc::new(new_storage);
            self.offset = 0;  // Reset offset since we have fresh storage
        }
        Ok(())
    }
    }
 /// Tensor core optimization analysis
 #[derive(Debug, Clone)]
 pub struct TensorCoreOptimization {
    pub can_use_tensor_cores: bool,
    pub recommended_precision: MixedPrecision,
    pub optimal_tile_size: Option<(usize, usize, usize)>,
    pub memory_layout: MemoryLayout,
 }
 /// Mixed precision options
 #[derive(Debug, Clone, Copy)]
 pub enum MixedPrecision {
    FP32,
    FP16,
    BF16,
    INT8,
 }
 /// Memory layout options
 #[derive(Debug, Clone, Copy)]
 pub enum MemoryLayout {
    RowMajor,
    ColumnMajor,
    TensorCore,
 }
 /// Check tensor core compatibility
 fn check_tensor_core_compatibility(shape: &[usize], dtype: crate::DType) -> bool {
    if shape.len() < 2 {
        return false;
    }
    let m = shape[shape.len() - 2];
    let n = shape[shape.len() - 1];
    (m % 16 == 0) && (n % 16 == 0)
 }
 /// Select optimal precision
 fn select_optimal_precision(shape: &[usize], _dtype: crate::DType) -> MixedPrecision {
    let total_elements: usize = shape.iter().product();
    if total_elements > 1_000_000 {
        MixedPrecision::FP16
    } else {
        MixedPrecision::FP32
    }
 }
 // Additional tensor operations for autodiff support
 impl Tensor {
    /// Create a scalar tensor
    pub fn scalar(value: f32, device: &Device) -> Self {
        Self::from_data(&[value], &[], device).unwrap_or_else(|_| {
            Self::zeros(&[], device).expect("Failed to create fallback zero scalar tensor")
        })
    }
    /// Create ones_like tensor
    pub fn ones_like(tensor: &Tensor) -> Self {
        Self::ones(tensor.shape().dims(), tensor.device()).unwrap_or_else(|_| {
            tensor.clone()
        })
    }
    /// Create zeros_like tensor
    pub fn zeros_like(tensor: &Tensor) -> Self {
        Self::zeros(tensor.shape().dims(), tensor.device()).unwrap_or_else(|_| {
            tensor.clone()
        })
    }
    /// Create tensor filled with specific value
    pub fn full<S: Into<Shape>>(shape: S, value: f32, device: &Device) -> Result<Self> {
        let shape = shape.into();
        let numel = shape.numel();
        let storage = Arc::new(Storage::full(numel, value, DType::F32, device)?);
        let strides = shape.default_strides();
        Ok(Tensor {
            storage,
            shape,
            strides,
            offset: 0,
            dtype: DType::F32,
            device: device.clone(),
            requires_grad: false,
            grad: Arc::new(Mutex::new(None)),
            node_id: None,
        })
    }
    /// Create tensor from vector
    pub fn from_vec(data: Vec<f32>, shape: &[usize], device: &Device) -> Self {
        Self::from_data(&data, shape, device).unwrap_or_else(|_| {
            Self::zeros(shape, device).expect("Failed to create fallback zero tensor")
        })
    }
    /// Element-wise sine
    pub fn sin(&self) -> Self {
        let data = self.to_cpu().unwrap_or_else(|_| vec![]);
        let result: Vec<f32> = data.iter().map(|x| x.sin()).collect();
        Self::from_vec(result, self.shape().dims(), self.device())
    }
    /// Element-wise cosine
    pub fn cos(&self) -> Self {
        let data = self.to_cpu().unwrap_or_else(|_| vec![]);
        let result: Vec<f32> = data.iter().map(|x| x.cos()).collect();
        Self::from_vec(result, self.shape().dims(), self.device())
    }
    /// Element-wise ReLU
    pub fn relu(&self) -> Self {
        let data = self.to_cpu().unwrap_or_else(|_| vec![]);
        let result: Vec<f32> = data.iter().map(|x| x.max(0.0)).collect();
        Self::from_vec(result, self.shape().dims(), self.device())
    }
    /// Compute mean
    pub fn mean(&self, dim: Option<usize>, _keepdim: bool) -> Self {
        let data = self.to_cpu().unwrap_or_else(|_| vec![]);
        if data.is_empty() {
            return Self::scalar(0.0, self.device());
        }
        if let Some(_d) = dim {
            // Mean along specific dimension - simplified for now
            let sum: f32 = data.iter().sum();
            let count = data.len() as f32;
            Self::scalar(sum / count, self.device())
        } else {
            // Mean of all elements
            let sum: f32 = data.iter().sum();
            let count = data.len() as f32;
            Self::scalar(sum / count, self.device())
        }
    }
    /// Multiply by scalar
    pub fn mul_scalar(&self, scalar: f32) -> Self {
        let data = self.to_cpu().unwrap_or_else(|_| vec![]);
        let result: Vec<f32> = data.iter().map(|x| x * scalar).collect();
        Self::from_vec(result, self.shape().dims(), self.device())
    }
    /// Subtract tensors
    pub fn sub(&self, other: &Tensor) -> Self {
        let self_data = self.to_cpu().unwrap_or_else(|_| vec![]);
        let other_data = other.to_cpu().unwrap_or_else(|_| vec![]);
        if self_data.len() != other_data.len() {
            return self.clone();
        }
        let result: Vec<f32> = self_data
            .iter()
            .zip(other_data.iter())
            .map(|(a, b)| a - b)
            .collect();
        Self::from_vec(result, self.shape().dims(), self.device())
    }
    /// Convert to vector (alias for to_cpu for compatibility)
    pub fn to_vec(&self) -> Vec<f32> {
        self.to_cpu().unwrap_or_else(|_| vec![])
    }
 }
@@ -1,234 +0,0 @@
 //! Knowledge distillation for model compression
 use crate::error::{CompressionError, Result};
 use rtx_tensor::{Device, Tensor};
 /// Distillation loss type
 #[derive(Debug, Clone, Copy)]
 pub enum DistillationLoss {
    KLDivergence,
    MSE,
    CrossEntropy,
 }
 /// Configuration for knowledge distillation
 #[derive(Debug, Clone)]
 pub struct DistillationConfig {
    pub temperature: f32,
    pub alpha: f32,
    pub beta: f32,
    pub loss_type: DistillationLoss,
    pub use_attention_transfer: bool,
    pub use_feature_matching: bool,
 }
 impl Default for DistillationConfig {
    fn default() -> Self {
        Self {
            temperature: 3.0,
            alpha: 0.7,
            beta: 0.3,
            loss_type: DistillationLoss::KLDivergence,
            use_attention_transfer: false,
            use_feature_matching: false,
        }
    }
 }
 /// Knowledge distiller for teacher-student compression
 #[derive(Debug, Clone)]
 pub struct KnowledgeDistiller {
    config: DistillationConfig,
    device: Device,
    schedule_steps: Option<usize>,
    initial_temp: f32,
    final_temp: f32,
 }
 impl KnowledgeDistiller {
    pub fn new(config: DistillationConfig, device: &Device) -> Result<Self> {
        Ok(Self {
            config,
            device: device.clone(),
            schedule_steps: None,
            initial_temp: config.temperature,
            final_temp: config.temperature,
        })
    }
    pub fn default(device: &Device) -> Result<Self> {
        Self::new(DistillationConfig::default(), device)
    }
    pub fn with_schedule(
        initial_temp: f32,
        final_temp: f32,
        steps: usize,
        device: &Device,
    ) -> Result<Self> {
        let config = DistillationConfig {
            temperature: initial_temp,
            ..Default::default()
        };
        Ok(Self {
            config,
            device: device.clone(),
            schedule_steps: Some(steps),
            initial_temp,
            final_temp,
        })
    }
    pub fn compute_loss(&self, teacher_logits: &Tensor, student_logits: &Tensor) -> Result<Tensor> {
        let teacher_soft = self.apply_temperature(teacher_logits)?;
        let student_soft = self.apply_temperature(student_logits)?;
        match self.config.loss_type {
            DistillationLoss::KLDivergence => self.kl_divergence(&teacher_soft, &student_soft),
            DistillationLoss::MSE => self.mse_loss(&teacher_soft, &student_soft),
            DistillationLoss::CrossEntropy => self.cross_entropy(&teacher_soft, &student_soft),
        }
    }
    pub fn apply_temperature(&self, logits: &Tensor) -> Result<Tensor> {
        logits.div_scalar(self.config.temperature)
    }
    pub fn generate_soft_targets(&self, logits: &Tensor) -> Result<Tensor> {
        let scaled = self.apply_temperature(logits)?;
        self.softmax(&scaled)
    }
    pub fn compute_combined_loss(
        &self,
        teacher_logits: &Tensor,
        student_logits: &Tensor,
        labels: &Tensor,
    ) -> Result<Tensor> {
        let distill_loss = self.compute_loss(teacher_logits, student_logits)?;
        let student_loss = self.cross_entropy(student_logits, labels)?;
        let weighted_distill = distill_loss.mul_scalar(self.config.alpha)?;
        let weighted_student = student_loss.mul_scalar(self.config.beta)?;
        weighted_distill.add(&weighted_student)
    }
    pub fn attention_transfer_loss(
        &self,
        teacher_attn: &Tensor,
        student_attn: &Tensor,
    ) -> Result<Tensor> {
        // Simplified: MSE between attention maps
        self.mse_loss(teacher_attn, student_attn)
    }
    pub fn feature_matching_loss(
        &self,
        teacher_features: &[Tensor],
        student_features: &[Tensor],
    ) -> Result<Tensor> {
        if teacher_features.len() != student_features.len() {
            return Err(CompressionError::CompressionFailed(
                "Feature lists must have same length".to_string()
            ));
        }
        let mut total_loss = Tensor::zeros(&[1], &self.device)?;
        for (t, s) in teacher_features.iter().zip(student_features.iter()) {
            let loss = self.mse_loss(t, s)?;
            total_loss = total_loss.add(&loss)?;
        }
        total_loss.div_scalar(teacher_features.len() as f32)
    }
    pub fn compute_compression_ratio(&self, teacher_params: usize, student_params: usize) -> f32 {
        teacher_params as f32 / student_params as f32
    }
    pub fn get_temperature_at_step(&self, step: usize) -> f32 {
        if let Some(total_steps) = self.schedule_steps {
            let progress = (step as f32) / (total_steps as f32).max(1.0);
            self.initial_temp + (self.final_temp - self.initial_temp) * progress
        } else {
            self.config.temperature
        }
    }
    pub fn compute_metrics(
        &self,
        teacher_preds: &Tensor,
        student_preds: &Tensor,
        labels: &Tensor,
    ) -> Result<DistillationMetrics> {
        let teacher_correct = self.count_correct(teacher_preds, labels)?;
        let student_correct = self.count_correct(student_preds, labels)?;
        let agreement = self.count_agreement(teacher_preds, student_preds)?;
        let total = labels.numel() as f32;
        Ok(DistillationMetrics {
            teacher_accuracy: teacher_correct / total,
            student_accuracy: student_correct / total,
            agreement_rate: agreement / total,
        })
    }
    // Helper methods
    fn kl_divergence(&self, p: &Tensor, q: &Tensor) -> Result<Tensor> {
        // Simplified KL divergence
        let log_p = p.log()?;
        let log_q = q.log()?;
        let diff = log_p.sub(&log_q)?;
        let kl = p.mul(&diff)?;
        kl.sum()?.div_scalar(p.numel() as f32)
    }
    fn mse_loss(&self, pred: &Tensor, target: &Tensor) -> Result<Tensor> {
        let diff = pred.sub(target)?;
        let squared = diff.mul(&diff)?;
        squared.sum()?.div_scalar(pred.numel() as f32)
    }
    fn cross_entropy(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
        // Simplified cross entropy
        let probs = self.softmax(logits)?;
        let log_probs = probs.log()?;
        log_probs.sum()?.mul_scalar(-1.0)?.div_scalar(logits.shape().dims()[0] as f32)
    }
    fn softmax(&self, logits: &Tensor) -> Result<Tensor> {
        let exp = logits.exp()?;
        let sum = exp.sum()?;
        exp.div(&sum)
    }
    fn count_correct(&self, preds: &Tensor, labels: &Tensor) -> Result<f32> {
        let pred_data = preds.to_vec()?;
        let label_data = labels.to_vec()?;
        let correct = pred_data.iter()
            .zip(label_data.iter())
            .filter(|(p, l)| (p - l).abs() < 1e-5)
            .count();
        Ok(correct as f32)
    }
    fn count_agreement(&self, preds1: &Tensor, preds2: &Tensor) -> Result<f32> {
        let data1 = preds1.to_vec()?;
        let data2 = preds2.to_vec()?;
        let agree = data1.iter()
            .zip(data2.iter())
            .filter(|(p1, p2)| (p1 - p2).abs() < 1e-5)
            .count();
        Ok(agree as f32)
    }
 }
 /// Metrics for distillation
 #[derive(Debug)]
 pub struct DistillationMetrics {
    pub teacher_accuracy: f32,
    pub student_accuracy: f32,
    pub agreement_rate: f32,
 }
@@ -1,128 +0,0 @@
 //! Structured pruning implementation
 use crate::error::{CompressionError, Result};
 use crate::pruning::config::{PruningConfig, PruningGranularity};
 use rtx_tensor::{Device, Tensor};
 /// Structured pruner for channel/filter pruning
 #[derive(Debug, Clone)]
 pub struct StructuredPruner {
    config: PruningConfig,
    device: Device,
 }
 impl StructuredPruner {
    /// Create a new structured pruner
    pub fn new(config: PruningConfig, device: &Device) -> Result<Self> {
        if !config.structured {
            return Err(CompressionError::CompressionFailed(
                "Config must have structured=true for StructuredPruner".to_string()
            ));
        }
        Ok(Self {
            config,
            device: device.clone(),
        })
    }
    /// Compute pruning mask for structured pruning
    pub fn compute_mask(&self, weights: &Tensor) -> Result<Tensor> {
        match self.config.granularity {
            PruningGranularity::Channel => self.compute_channel_mask(weights),
            PruningGranularity::Filter => self.compute_filter_mask(weights),
            _ => Err(CompressionError::CompressionFailed(
                "Unsupported granularity for structured pruning".to_string()
            )),
        }
    }
    /// Apply mask to weights
    pub fn apply_mask(&self, weights: &Tensor, mask: &Tensor) -> Result<Tensor> {
        weights.mul(mask)
    }
    /// Prune weights directly
    pub fn prune(&self, weights: &Tensor) -> Result<Tensor> {
        let mask = self.compute_mask(weights)?;
        self.apply_mask(weights, &mask)
    }
    /// Get indices of pruned channels/filters
    pub fn get_pruned_indices(&self, weights: &Tensor) -> Result<Vec<usize>> {
        let shape = weights.shape().dims();
        let num_filters = shape[0];
        let num_to_prune = (num_filters as f32 * self.config.sparsity) as usize;
        // Compute L2 norm of each filter
        let data = weights.to_vec()?;
        let filter_size = data.len() / num_filters;
        let mut norms: Vec<(f32, usize)> = Vec::new();
        for i in 0..num_filters {
            let start = i * filter_size;
            let end = start + filter_size;
            let norm: f32 = data[start..end].iter()
                .map(|x| x * x)
                .sum::<f32>()
                .sqrt();
            norms.push((norm, i));
        }
        // Sort by norm and select smallest
        norms.sort_by(|a, b| a.0.total_cmp(&b.0));
        Ok(norms.iter().take(num_to_prune).map(|(_, idx)| *idx).collect())
    }
    fn compute_channel_mask(&self, weights: &Tensor) -> Result<Tensor> {
        let shape = weights.shape().dims();
        if shape.len() != 4 {
            return Err(CompressionError::CompressionFailed(
                "Channel pruning requires 4D tensor [out, in, h, w]".to_string()
            ));
        }
        let out_channels = shape[0];
        let channel_size = shape[1] * shape[2] * shape[3];
        let num_to_prune = (out_channels as f32 * self.config.sparsity) as usize;
        // Compute channel importance (L2 norm)
        let data = weights.to_vec()?;
        let mut channel_norms: Vec<(f32, usize)> = Vec::new();
        for ch in 0..out_channels {
            let start = ch * channel_size;
            let end = start + channel_size;
            let norm: f32 = data[start..end].iter()
                .map(|x| x * x)
                .sum::<f32>()
                .sqrt();
            channel_norms.push((norm, ch));
        }
        // Sort and select channels to prune
        channel_norms.sort_by(|a, b| a.0.total_cmp(&b.0));
        let pruned_channels: Vec<usize> = channel_norms.iter()
            .take(num_to_prune)
            .map(|(_, idx)| *idx)
            .collect();
        // Create mask
        let mut mask = vec![1.0f32; data.len()];
        for &ch in &pruned_channels {
            let start = ch * channel_size;
            let end = start + channel_size;
            for i in start..end {
                mask[i] = 0.0;
            }
        }
        Tensor::from_data(mask, shape.to_vec(), &self.device)
    }
    fn compute_filter_mask(&self, weights: &Tensor) -> Result<Tensor> {
        // Similar to channel mask but for entire filters
        self.compute_channel_mask(weights)
    }
 }
@@ -1,687 +0,0 @@
 //! Parallelism strategies for distributed training
 //!
 //! This module implements various parallelism approaches including:
 //! - Data Parallel (DP)
 //! - Tensor Parallel (TP)
 //! - Pipeline Parallel (PP)
 //! - Fully Sharded Data Parallel (FSDP/ZeRO)
 use crate::comm::CommunicationPrimitive;
 use crate::error::{DistributedError, Result};
 use crate::group::ProcessGroup;
 use crate::tensor_ext::{TensorExt, TensorShapeExt};
 use rtx_tensor::Tensor;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 /// Data Parallel implementation
 pub struct DataParallel {
    /// Process group for data parallel communication
    pub process_group: ProcessGroup,
    /// Gradient accumulation steps
    pub accumulation_steps: usize,
    /// Current accumulation step
    current_step: usize,
    /// Whether gradients are currently being accumulated
    accumulating: bool,
 }
 impl DataParallel {
    /// Create new data parallel instance
    pub fn new(process_group: ProcessGroup, accumulation_steps: usize) -> Self {
        Self {
            process_group,
            accumulation_steps: accumulation_steps.max(1),
            current_step: 0,
            accumulating: false,
        }
    }
    /// Start gradient accumulation
    pub fn start_accumulation(&mut self) {
        self.accumulating = true;
        self.current_step = 0;
    }
    /// Accumulate gradients (called after backward pass)
    pub async fn accumulate_gradients(&mut self, gradients: &mut [Tensor]) -> Result<bool> {
        if !self.accumulating {
            return Err(DistributedError::parallelism(
            "DataParallel",
            "accumulation not started"
            ));
        }
        self.current_step += 1;
        // If we've accumulated enough steps, average gradients
        if self.current_step >= self.accumulation_steps {
            self.average_gradients(gradients).await?;
            self.accumulating = false;
            self.current_step = 0;
            Ok(true) // Ready for optimizer step
            } else {
                Ok(false) // Continue accumulating
            }
        }
        /// Average gradients across all processes
        async fn average_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
            use crate::comm::ReduceOp;
            for gradient in gradients.iter_mut() {
                // AllReduce sum
                self.process_group.allreduce(gradient, ReduceOp::Sum).await?;
                // Divide by world size and accumulation steps
                let scale = 1.0 / (self.process_group.world_size() as f32 * self.accumulation_steps as f32);
                gradient.mul_scalar_(scale)?;
            }
            Ok(())
        }
        /// Get current accumulation progress
        pub fn accumulation_progress(&self) -> f32 {
            if self.accumulation_steps == 0 {
                1.0
                } else {
                    self.current_step as f32 / self.accumulation_steps as f32
                }
            }
        }
        /// Tensor Parallel implementation
        pub struct TensorParallel {
            /// Process group for tensor parallel communication
            pub process_group: ProcessGroup,
            /// Tensor sharding configuration
            pub sharding_config: TensorShardingConfig,
        }
        /// Configuration for tensor sharding
        #[derive(Debug, Clone, Serialize, Deserialize)]
        pub struct TensorShardingConfig {
            /// Dimension along which to shard tensors
            pub shard_dim: usize,
            /// Whether to shard attention weights
            pub shard_attention: bool,
            /// Whether to shard feed-forward weights
            pub shard_feedforward: bool,
            /// Overlap communication with computation
            pub overlap_comm: bool,
        }
        impl Default for TensorShardingConfig {
            fn default() -> Self {
                Self {
                    shard_dim: 0, // Shard along first dimension
                    shard_attention: true,
                    shard_feedforward: true,
                    overlap_comm: true,
                }
            }
        }
        impl TensorParallel {
            /// Create new tensor parallel instance
            pub fn new(process_group: ProcessGroup, sharding_config: TensorShardingConfig) -> Self {
                Self {
                    process_group,
                    sharding_config,
                }
            }
            /// Shard a tensor across processes
            pub async fn shard_tensor(&self, tensor: &Tensor) -> Result<Tensor> {
                let world_size = self.process_group.world_size() as usize;
                let shard_dim = self.sharding_config.shard_dim;
                // Check if tensor can be sharded
                if shard_dim >= tensor.shape().dims().len() {
                    return Err(DistributedError::tensor(
                    "shard dimension exceeds tensor dimensions"
                    ));
                }
                let dim_size = tensor.shape().dims()[shard_dim];
                if dim_size % world_size != 0 {
                    return Err(DistributedError::tensor(
                    format!("tensor dimension {} not divisible by world size {}", dim_size, world_size)
                    ));
                }
                // Calculate shard size
                let shard_size = dim_size / world_size;
                let my_rank = self.process_group.rank() as usize;
                let start_idx = my_rank * shard_size;
                let end_idx = start_idx + shard_size;
                // Create sharded tensor view
                tensor.slice(shard_dim, start_idx, end_idx)
            }
            /// Gather tensor shards from all processes
            pub async fn gather_tensor_shards(&self, shard: &Tensor) -> Result<Tensor> {
                use crate::comm::AllGatherOutput;
                let gathered = self.process_group.allgather(shard).await?;
                match gathered {
                    AllGatherOutput::Tensor(tensor) => Ok(tensor),
                    AllGatherOutput::TensorList(tensors) => {
                        // Concatenate tensors along shard dimension
                        Tensor::cat(&tensors, self.sharding_config.shard_dim)
                    }
                }
            }
        }
        /// Pipeline Parallel implementation
        pub struct PipelineParallel {
            /// Process group for pipeline parallel communication
            pub process_group: ProcessGroup,
            /// Pipeline stage configuration
            pub stage_config: PipelineStageConfig,
            /// Number of microbatches for pipeline
            pub num_microbatches: usize,
        }
        /// Configuration for pipeline stages
        #[derive(Debug, Clone, Serialize, Deserialize)]
        pub struct PipelineStageConfig {
            /// Stage ID (0-indexed)
            pub stage_id: usize,
            /// Total number of stages
            pub num_stages: usize,
            /// Layers assigned to this stage
            pub layer_range: (usize, usize), // (start, end)
            /// Use asynchronous pipeline schedule
            pub async_schedule: bool,
        }
        impl PipelineParallel {
            /// Create new pipeline parallel instance
            pub fn new(
            process_group: ProcessGroup,
            stage_config: PipelineStageConfig,
            num_microbatches: usize
            ) -> Self {
                Self {
                    process_group,
                    stage_config,
                    num_microbatches,
                }
            }
            /// Execute forward pass for pipeline stage
            pub async fn forward_stage(&self, input: &Tensor, stage_layers: &[Box<dyn LayerTrait>]) -> Result<Tensor> {
                let mut output = input.clone();
                // Apply layers for this stage
                for layer in stage_layers {
                    output = layer.forward(&output)?;
                }
                // Send to next stage if not the last stage
                if self.stage_config.stage_id < self.stage_config.num_stages - 1 {
                    let next_rank = self.process_group.rank() + 1;
                    self.process_group.send(&output, next_rank).await?;
                }
                Ok(output)
            }
            /// Execute backward pass for pipeline stage
            pub async fn backward_stage(&self, grad_output: &Tensor, stage_layers: &[Box<dyn LayerTrait>]) -> Result<Tensor> {
                let mut grad_input = grad_output.clone();
                // Apply backward pass for layers in reverse order
                for layer in stage_layers.iter().rev() {
                    grad_input = layer.backward(&grad_input)?;
                }
                // Send to previous stage if not the first stage
                if self.stage_config.stage_id > 0 {
                    let prev_rank = self.process_group.rank() - 1;
                    self.process_group.send(&grad_input, prev_rank).await?;
                }
                Ok(grad_input)
            }
            /// Check if this is the first stage
            pub fn is_first_stage(&self) -> bool {
                self.stage_config.stage_id == 0
            }
            /// Check if this is the last stage
            pub fn is_last_stage(&self) -> bool {
                self.stage_config.stage_id == self.stage_config.num_stages - 1
            }
        }
        /// Placeholder trait for neural network layers
        pub trait LayerTrait {
            fn forward(&self, input: &Tensor) -> Result<Tensor>;
            fn backward(&self, grad_output: &Tensor) -> Result<Tensor>;
        }
        /// Fully Sharded Data Parallel (FSDP) implementation
        pub struct Fsdp {
            /// Process group for FSDP communication
            pub process_group: ProcessGroup,
            /// FSDP configuration
            pub config: FsdpConfig,
            /// Sharded parameters
            sharded_params: HashMap<String, ShardedParameter>,
            /// Memory usage statistics
            memory_stats: FsdpMemoryStats,
        }
        /// FSDP configuration
        #[derive(Debug, Clone, Serialize, Deserialize)]
        pub struct FsdpConfig {
            /// Sharding strategy
            pub sharding_strategy: ShardingStrategy,
            /// Minimum parameter size to shard (in elements)
            pub min_param_size: usize,
            /// CPU offloading enabled
            pub cpu_offload: bool,
            /// Mixed precision configuration
            pub mixed_precision: bool,
            /// Flatten parameters for sharding
            pub flatten_parameters: bool,
        }
        /// Sharding strategies for FSDP
        #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
        pub enum ShardingStrategy {
            /// Full sharding (ZeRO-3)
            FullShard,
            /// Shard gradients only (ZeRO-2)
            ShardGradOp,
            /// No sharding, replicate parameters
            NoShard,
        }
        /// Sharded parameter representation
        #[derive(Debug)]
        pub struct ShardedParameter {
            /// Parameter name/identifier
            pub name: String,
            /// Local shard of the parameter
            pub local_shard: Tensor,
            /// Full parameter shape
            pub full_shape: Vec<usize>,
            /// Shard metadata
            pub shard_metadata: ShardMetadata,
        }
        /// Metadata for parameter sharding
        #[derive(Debug, Clone, Serialize, Deserialize)]
        pub struct ShardMetadata {
            /// Rank that owns this shard
            pub owner_rank: i32,
            /// Start index in flattened parameter
            pub start_idx: usize,
            /// End index in flattened parameter
            pub end_idx: usize,
            /// Original parameter offset
            pub param_offset: usize,
        }
        /// Memory usage statistics for FSDP
        #[derive(Debug, Default)]
        pub struct FsdpMemoryStats {
            /// Total parameter memory (MB)
            pub total_param_memory_mb: f32,
            /// Local shard memory (MB)
            pub local_shard_memory_mb: f32,
            /// Peak memory during all-gather (MB)
            pub peak_allgather_memory_mb: f32,
            /// Memory saved compared to non-sharded (MB)
            pub memory_saved_mb: f32,
        }
        impl Default for FsdpConfig {
            fn default() -> Self {
                Self {
                    sharding_strategy: ShardingStrategy::FullShard,
                    min_param_size: 1000, // Only shard parameters with >1K elements
                    cpu_offload: false,
                    mixed_precision: true,
                    flatten_parameters: true,
                }
            }
        }
        impl Fsdp {
            /// Create new FSDP instance with default config
            pub fn new(process_group: ProcessGroup) -> Result<Self> {
                let config = FsdpConfig::default();
                Ok(Self {
                    process_group,
                    config,
                    sharded_params: HashMap::new(),
                    memory_stats: FsdpMemoryStats::default(),
                })
            }
            /// Create new FSDP instance with custom config
            pub fn new_with_config(process_group: ProcessGroup, config: FsdpConfig) -> Self {
                Self {
                    process_group,
                    config,
                    sharded_params: HashMap::new(),
                    memory_stats: FsdpMemoryStats::default(),
                }
            }
            /// Shard model parameters across the process group
            pub fn shard_parameters(&mut self, parameters: &Tensor) -> Result<Tensor> {
                let world_size = self.process_group.world_size();
                let rank = self.process_group.rank();
                if world_size == 1 {
                    // No sharding needed for single GPU
                    return Ok(parameters.clone());
                }
                // Calculate shard size
                let total_elements = parameters.element_count();
                let elements_per_shard = (total_elements + world_size - 1) / world_size; // Round up
                let start_idx = rank * elements_per_shard;
                let end_idx = (start_idx + elements_per_shard).min(total_elements);
                let shard_size = end_idx - start_idx;
                // Create local shard (simplified - in real implementation would slice the actual tensor)
                let shard_shape = rtx_tensor::TensorShape::new(vec![shard_size]);
                let local_shard = Tensor::zeros(shard_shape);
                // Update memory statistics
                let total_memory_mb = (total_elements * 4) as f32 / (1024.0 * 1024.0); // f32 = 4 bytes
                let shard_memory_mb = (shard_size * 4) as f32 / (1024.0 * 1024.0);
                self.memory_stats.total_param_memory_mb = total_memory_mb;
                self.memory_stats.local_shard_memory_mb = shard_memory_mb;
                self.memory_stats.memory_saved_mb = total_memory_mb - shard_memory_mb;
                tracing::debug!("FSDP sharded {:.2}MB -> {:.2}MB ({:.1}% reduction)",
                total_memory_mb, shard_memory_mb,
                (self.memory_stats.memory_saved_mb / total_memory_mb) * 100.0);
                Ok(local_shard)
            }
            /// Synchronize gradients across all processes
            pub fn sync_gradients(&self, gradients: &mut Tensor) -> Result<()> {
                use crate::comm::ReduceOp;
                // Perform AllReduce to sum gradients across all processes
                self.process_group.all_reduce(gradients, ReduceOp::Sum)?;
                // Average the gradients
                let world_size = self.process_group.world_size() as f32;
                *gradients = gradients.clone() / world_size;
                Ok(())
            }
            /// Get current memory statistics
            pub fn memory_stats(&self) -> &FsdpMemoryStats {
                &self.memory_stats
            }
        }
        let world_size = self.process_group.world_size() as usize;
        let my_rank = self.process_group.rank() as usize;
        for (param_idx, param) in parameters.iter().enumerate() {
            let param_name = format!("param_{}", param_idx);
            // Skip small parameters if configured
            let param_size = param.shape().dims().iter().product::<usize>();
            if param_size < self.config.min_param_size {
                continue;
            }
            // Flatten parameter if configured
            let flattened_param = if self.config.flatten_parameters {
                param.flatten()?
                } else {
                    param.clone()
                };
                // Calculate shard boundaries
                let total_elements = flattened_param.shape().dims().iter().product::<usize>();
                let elements_per_shard = (total_elements + world_size - 1) / world_size; // Ceiling division
                let start_idx = my_rank * elements_per_shard;
                let end_idx = (start_idx + elements_per_shard).min(total_elements);
                // Create local shard
                let local_shard = if start_idx < total_elements {
                    flattened_param.slice(0, start_idx, end_idx)?
                    } else {
                        // Empty shard for ranks with no data
                        Tensor::zeros(rtx_tensor::TensorShape::new(vec![0]))?
                    };
                    // Create shard metadata
                    let shard_metadata = ShardMetadata {
                        owner_rank: my_rank as i32,
                        start_idx,
                        end_idx,
                        param_offset: param_idx,
                    };
                    // Store sharded parameter
                    let sharded_param = ShardedParameter {
                        name: param_name.clone(),
                        local_shard,
                        full_shape: param.shape().dims().to_vec(),
                        shard_metadata,
                    };
                    self.sharded_params.insert(param_name, sharded_param);
                }
                // Update memory statistics
                self.update_memory_stats(parameters).await?;
                Ok(())
            }
            /// All-gather parameters for forward pass
            pub async fn allgather_parameters(&self, param_name: &str) -> Result<Tensor> {
                let sharded_param = self.sharded_params.get(param_name)
                .ok_or_else(|| DistributedError::tensor(
                format!("parameter {} not found", param_name)
                ))?;
                // All-gather the sharded parameter
                use crate::comm::AllGatherOutput;
                let gathered = self.process_group.allgather(&sharded_param.local_shard).await?;
                match gathered {
                    AllGatherOutput::Tensor(tensor) => {
                        // Reshape to original shape if flattened
                        if self.config.flatten_parameters && sharded_param.full_shape.len() > 1 {
                            tensor.reshape(&sharded_param.full_shape)
                            } else {
                                Ok(tensor)
                            }
                        },
                        AllGatherOutput::TensorList(_) => {
                            Err(DistributedError::tensor("unexpected tensor list in all-gather"))
                        }
                    }
                }
                /// Reduce-scatter gradients after backward pass
                pub async fn reduce_scatter_gradients(&self, param_name: &str, full_gradient: &Tensor) -> Result<()> {
                    let _sharded_param = self.sharded_params.get(param_name)
                    .ok_or_else(|| DistributedError::tensor(
                    format!("parameter {} not found", param_name)
                    ))?;
                    // Flatten gradient if needed
                    let flattened_grad = if self.config.flatten_parameters {
                        full_gradient.flatten()?
                        } else {
                            full_gradient.clone()
                        };
                        // Reduce-scatter the gradient
                        use crate::comm::ReduceOp;
                        let _local_grad_shard = self.process_group
                        .reduce_scatter(&flattened_grad, ReduceOp::Sum).await?;
                        // In a real implementation, would update the local parameter shard here
                        Ok(())
                    }
                    /// Update memory usage statistics
                    async fn update_memory_stats(&mut self, original_params: &[Tensor]) -> Result<()> {
                        // Calculate total parameter memory
                        let mut total_param_memory = 0.0;
                        for param in original_params {
                            let param_size = param.shape().dims().iter().product::<usize>();
                            total_param_memory += param_size as f32 * 4.0; // Assuming f32
                        }
                        self.memory_stats.total_param_memory_mb = total_param_memory / (1024.0 * 1024.0);
                        // Calculate local shard memory
                        let mut local_shard_memory = 0.0;
                        for sharded_param in self.sharded_params.values() {
                            let shard_size = sharded_param.local_shard.shape().dims().iter().product::<usize>();
                            local_shard_memory += shard_size as f32 * 4.0; // Assuming f32
                        }
                        self.memory_stats.local_shard_memory_mb = local_shard_memory / (1024.0 * 1024.0);
                        // Calculate memory savings
                        self.memory_stats.memory_saved_mb =
                        self.memory_stats.total_param_memory_mb - self.memory_stats.local_shard_memory_mb;
                        // Peak all-gather memory is approximately total parameter memory
                        self.memory_stats.peak_allgather_memory_mb = self.memory_stats.total_param_memory_mb;
                        Ok(())
                    }
                    /// Get memory usage statistics
                    pub fn memory_stats(&self) -> &FsdpMemoryStats {
                        &self.memory_stats
                    }
                    /// Calculate memory reduction percentage
                    pub fn memory_reduction_percent(&self) -> f32 {
                        if self.memory_stats.total_param_memory_mb > 0.0 {
                            (self.memory_stats.memory_saved_mb / self.memory_stats.total_param_memory_mb) * 100.0
                            } else {
                                0.0
                            }
                        }
                    }
                    #[cfg(test)]
                    mod tests {
                        use super::*;
                        use crate::{Backend, BackendConfig};
                        use rtx_tensor::TensorShape;
                        #[tokio::test]
                        async fn test_data_parallel() {
                            let config = BackendConfig::cpu();
                            let pg = ProcessGroup::new(Backend::Cpu, 2, 0, config).await.unwrap();
                            let mut dp = DataParallel::new(pg, 2);
                            dp.start_accumulation();
                            assert_eq!(dp.accumulation_progress(), 0.0);
                            // Create dummy gradients
                            let shape = TensorShape::new(vec![4]);
                            let mut gradients = vec![Tensor::ones(shape).unwrap()];
                            // First accumulation
                            let ready = dp.accumulate_gradients(&mut gradients).await.unwrap();
                            assert!(!ready);
                            assert_eq!(dp.accumulation_progress(), 0.5);
                            // Second accumulation should complete
                            let ready = dp.accumulate_gradients(&mut gradients).await.unwrap();
                            assert!(ready);
                            assert_eq!(dp.accumulation_progress(), 0.0);
                        }
                        #[tokio::test]
                        async fn test_tensor_parallel() {
                            let config = BackendConfig::cpu();
                            let pg = ProcessGroup::new(Backend::Cpu, 2, 0, config).await.unwrap();
                            let sharding_config = TensorShardingConfig::default();
                            let tp = TensorParallel::new(pg, sharding_config);
                            // Create tensor that can be sharded (dimension 0 divisible by world size)
                            let shape = TensorShape::new(vec![4, 3]); // 4 is divisible by 2
                            let tensor = Tensor::ones(shape).unwrap();
                            let shard = tp.shard_tensor(&tensor).await.unwrap();
                            assert_eq!(shard.shape().dims()[0], 2); // 4/2 = 2
                            assert_eq!(shard.shape().dims()[1], 3); // Unchanged
                        }
                        #[tokio::test]
                        async fn test_pipeline_parallel() {
                            let config = BackendConfig::cpu();
                            let pg = ProcessGroup::new(Backend::Cpu, 2, 0, config).await.unwrap();
                            let stage_config = PipelineStageConfig {
                                stage_id: 0,
                                num_stages: 2,
                                layer_range: (0, 2),
                                async_schedule: false,
                            };
                            let pp = PipelineParallel::new(pg, stage_config, 4);
                            assert!(pp.is_first_stage());
                            assert!(!pp.is_last_stage());
                        }
                        #[tokio::test]
                        async fn test_fsdp_sharding() {
                            let config = BackendConfig::cpu();
                            let pg = ProcessGroup::new(Backend::Cpu, 2, 0, config).await.unwrap();
                            let fsdp_config = FsdpConfig::default();
                            let mut fsdp = Fsdp::new(pg, fsdp_config);
                            // Create parameters
                            let shape1 = TensorShape::new(vec![2000]); // Large enough to shard
                            let shape2 = TensorShape::new(vec![100]);  // Too small to shard
                            let params = vec![
                            Tensor::ones(shape1).unwrap(),
                            Tensor::ones(shape2).unwrap(),
                            ];
                            fsdp.shard_parameters(&params).await.unwrap();
                            // Should only shard the first parameter (large enough)
                            assert_eq!(fsdp.sharded_params.len(), 1);
                            let memory_reduction = fsdp.memory_reduction_percent();
                            assert!(memory_reduction > 0.0);
                            assert!(memory_reduction < 100.0);
                        }
                        #[test]
                        fn test_sharding_strategy() {
                            assert_eq!(ShardingStrategy::FullShard, ShardingStrategy::FullShard);
                            assert_ne!(ShardingStrategy::FullShard, ShardingStrategy::ShardGradOp);
                        }
                        #[test]
                        fn test_fsdp_config_default() {
                            let config = FsdpConfig::default();
                            assert_eq!(config.sharding_strategy, ShardingStrategy::FullShard);
                            assert_eq!(config.min_param_size, 1000);
                            assert!(config.mixed_precision);
                        }
                    }
@@ -1,529 +0,0 @@
 //! Core Flash Attention implementation
 use crate::{
    config::FlashAttentionConfig,
    error::{FlashError, FlashResult},
    memory::SRAMManager,
    kernels::FlashCudaKernels,
    FlashOutput, FlashGradOutput, FlashStats,
 };
 use rtx_tensor::Tensor;
 use rtx_runtime::{CudaBackend, Stream};
 use std::sync::Arc;
 use parking_lot::RwLock;
 use async_trait::async_trait;
 use tracing::{info, debug, warn, error};
 /// Flash Attention backend trait for different implementations
 #[async_trait]
 pub trait FlashAttentionBackend {
    /// Forward pass of Flash Attention
    async fn forward(
    &self,
    q: &Tensor,      // Query tensor [batch, heads, seq_len, head_dim]
    k: &Tensor,      // Key tensor [batch, heads, seq_len, head_dim]
    v: &Tensor,      // Value tensor [batch, heads, seq_len, head_dim]
    causal: bool,
    softmax_scale: f32,
    ) -> FlashResult<FlashOutput>;
    /// Backward pass of Flash Attention
    async fn backward(
    &self,
    dout: &Tensor,   // Gradient w.r.t output [batch, heads, seq_len, head_dim]
    q: &Tensor,      // Query tensor
    k: &Tensor,      // Key tensor
    v: &Tensor,      // Value tensor
    output: &Tensor, // Forward output
    lse: &Tensor,    // Log-sum-exp from forward pass
    causal: bool,
    softmax_scale: f32,
    ) -> FlashResult<FlashGradOutput>;
    /// Get backend name for identification
    fn name(&self) -> &str;
    /// Check if backend supports the given configuration
    fn supports_config(&self, config: &FlashAttentionConfig) -> bool;
    /// Get optimal configuration for this backend
    fn optimize_config(&self, config: FlashAttentionConfig) -> FlashResult<FlashAttentionConfig>;
 }
 /// Main Flash Attention implementation
 pub struct FlashAttention {
    config: FlashAttentionConfig,
    cuda_backend: Arc<CudaBackend>,
    kernels: Arc<FlashCudaKernels>,
    sram_manager: Arc<SRAMManager>,
    streams: Arc<RwLock<Vec<Stream>>>,
    stats: Arc<RwLock<FlashStats>>,
 }
 impl FlashAttention {
    /// Create a new Flash Attention instance
    pub fn new(config: FlashAttentionConfig) -> FlashResult<Self> {
        // Validate configuration
        config.validate()?;
        info!("Initializing Flash Attention with config: {:?}", config);
        // Initialize CUDA backend
        let cuda_backend = Arc::new(CudaBackend::new(config.device_id)
        .map_err(|e| FlashError::backend_init(format!("Failed to initialize CUDA backend: {}", e)))?);
        // Initialize CUDA kernels
        let kernels = Arc::new(FlashCudaKernels::new(&cuda_backend, &config)?);
        // Initialize SRAM manager
        let sram_manager = Arc::new(SRAMManager::new(&config, &cuda_backend)?);
        // Initialize CUDA streams
        let mut streams = Vec::new();
        for i in 0..config.backend_config.cuda.stream_pool_size {
            let stream = cuda_backend.create_stream()
            .map_err(|e| FlashError::backend_init(format!("Failed to create stream {}: {}", i, e)))?;
            streams.push(stream);
        }
        let stats = Arc::new(RwLock::new(FlashStats {
            forward_time_us: 0,
            backward_time_us: 0,
            memory_usage: 0,
            sram_efficiency: 0.0,
            kernel_occupancy: 0.0,
        }));
        Ok(Self {
            config,
            cuda_backend,
            kernels,
            sram_manager,
            streams: Arc::new(RwLock::new(streams)),
            stats,
        })
    }
    /// Get configuration
    pub fn config(&self) -> &FlashAttentionConfig {
        &self.config
    }
    /// Get execution statistics
    pub fn stats(&self) -> FlashStats {
        self.stats.read().clone()
    }
    /// Validate input tensors
    fn validate_inputs(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> FlashResult<()> {
        // Check tensor devices
        if q.device() != k.device() || k.device() != v.device() {
            return Err(FlashError::config("All tensors must be on the same device"));
        }
        // Check tensor dtypes
        if q.dtype() != k.dtype() || k.dtype() != v.dtype() {
            return Err(FlashError::config("All tensors must have the same dtype"));
        }
        // Check tensor shapes
        let q_shape = q.shape();
        let k_shape = k.shape();
        let v_shape = v.shape();
        if q_shape.len() != 4 || k_shape.len() != 4 || v_shape.len() != 4 {
            return Err(FlashError::config("Input tensors must be 4D [batch, heads, seq_len, head_dim]"));
        }
        let [batch_q, heads_q, seq_q, dim_q] = q_shape[..] else {
            return Err(FlashError::config("Invalid Q tensor shape"));
        };
        let [batch_k, heads_k, seq_k, dim_k] = k_shape[..] else {
            return Err(FlashError::config("Invalid K tensor shape"));
        };
        let [batch_v, heads_v, seq_v, dim_v] = v_shape[..] else {
            return Err(FlashError::config("Invalid V tensor shape"));
        };
        if batch_q != batch_k || batch_k != batch_v {
            return Err(FlashError::shape_mismatch(
            vec![batch_q, heads_q, seq_q, dim_q],
            vec![batch_k, heads_k, seq_k, dim_k],
            ));
        }
        if heads_q != heads_k || heads_k != heads_v {
            return Err(FlashError::shape_mismatch(
            vec![batch_q, heads_q, seq_q, dim_q],
            vec![batch_k, heads_k, seq_k, dim_k],
            ));
        }
        if seq_q != seq_k || seq_k != seq_v {
            return Err(FlashError::shape_mismatch(
            vec![batch_q, heads_q, seq_q, dim_q],
            vec![batch_k, heads_k, seq_k, dim_k],
            ));
        }
        if dim_q != dim_k || dim_k != dim_v {
            return Err(FlashError::shape_mismatch(
            vec![batch_q, heads_q, seq_q, dim_q],
            vec![batch_k, heads_k, seq_k, dim_k],
            ));
        }
        // Check dimensions match configuration
        if heads_q != self.config.num_heads {
            return Err(FlashError::config(format!(
            "Number of heads mismatch: expected {}, got {}",
            self.config.num_heads, heads_q
            )));
        }
        if dim_q != self.config.head_dim {
            return Err(FlashError::config(format!(
            "Head dimension mismatch: expected {}, got {}",
            self.config.head_dim, dim_q
            )));
        }
        // Check sequence length constraints
        if seq_q > self.config.max_seq_len {
            return Err(FlashError::config(format!(
            "Sequence length {} exceeds maximum {}",
            seq_q, self.config.max_seq_len
            )));
        }
        Ok(())
    }
    /// Get next available stream
    fn get_stream(&self) -> Stream {
        let streams = self.streams.read();
        // Simple round-robin selection for now
        // In the future, we could implement more sophisticated stream selection
        streams[0].clone()
    }
    /// Update statistics
    fn update_stats(&self, forward_time: u64, memory_usage: usize, sram_efficiency: f32, occupancy: f32) {
        let mut stats = self.stats.write();
        stats.forward_time_us = forward_time;
        stats.memory_usage = memory_usage;
        stats.sram_efficiency = sram_efficiency;
        stats.kernel_occupancy = occupancy;
    }
 }
 #[async_trait]
 impl FlashAttentionBackend for FlashAttention {
    async fn forward(
    &self,
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    causal: bool,
    softmax_scale: f32,
    ) -> FlashResult<FlashOutput> {
        debug!("Starting Flash Attention forward pass");
        let start_time = std::time::Instant::now();
        // Validate inputs
        self.validate_inputs(q, k, v)?;
        let batch_size = q.shape()[0];
        let seq_len = q.shape()[2];
        // Get stream for execution
        let stream = self.get_stream();
        // Allocate output tensors
        let output_shape = q.shape().to_vec();
        let lse_shape = vec![batch_size, self.config.num_heads, seq_len];
        let output = Tensor::zeros(&output_shape, q.dtype(), q.device())?;
        let lse = Tensor::zeros(&lse_shape, q.dtype(), q.device())?;
        // Execute Flash Attention kernel
        let kernel_result = self.kernels.flash_attention_forward(
        q, k, v, &output, &lse,
        causal, softmax_scale, &stream
        ).await?;
        let elapsed = start_time.elapsed();
        let forward_time_us = elapsed.as_micros() as u64;
        // Calculate memory usage
        let memory_usage = self.config.estimate_memory_usage(batch_size, seq_len);
        // Get SRAM efficiency from manager
        let sram_efficiency = self.sram_manager.get_efficiency();
        // Get kernel occupancy from result
        let kernel_occupancy = kernel_result.occupancy;
        // Update statistics
        self.update_stats(forward_time_us, memory_usage, sram_efficiency, kernel_occupancy);
        let stats = FlashStats {
            forward_time_us,
            backward_time_us: 0,
            memory_usage,
            sram_efficiency,
            kernel_occupancy,
        };
        info!("Flash Attention forward completed in {}μs", forward_time_us);
        Ok(FlashOutput {
            output,
            lse,
            stats,
        })
    }
    async fn backward(
    &self,
    dout: &Tensor,
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    output: &Tensor,
    lse: &Tensor,
    causal: bool,
    softmax_scale: f32,
    ) -> FlashResult<FlashGradOutput> {
        debug!("Starting Flash Attention backward pass");
        let start_time = std::time::Instant::now();
        // Validate inputs
        self.validate_inputs(q, k, v)?;
        self.validate_inputs(dout, q, k)?; // dout should match q shape
        let batch_size = q.shape()[0];
        let seq_len = q.shape()[2];
        // Get stream for execution
        let stream = self.get_stream();
        // Allocate gradient tensors
        let dq = Tensor::zeros(q.shape(), q.dtype(), q.device())?;
        let dk = Tensor::zeros(k.shape(), k.dtype(), k.device())?;
        let dv = Tensor::zeros(v.shape(), v.dtype(), v.device())?;
        // Execute Flash Attention backward kernel
        let kernel_result = self.kernels.flash_attention_backward(
        dout, q, k, v, output, lse,
        &dq, &dk, &dv,
        causal, softmax_scale, &stream
        ).await?;
        let elapsed = start_time.elapsed();
        let backward_time_us = elapsed.as_micros() as u64;
        // Calculate memory usage
        let memory_usage = self.config.estimate_memory_usage(batch_size, seq_len);
        // Get SRAM efficiency from manager
        let sram_efficiency = self.sram_manager.get_efficiency();
        // Get kernel occupancy from result
        let kernel_occupancy = kernel_result.occupancy;
        let stats = FlashStats {
            forward_time_us: 0,
            backward_time_us,
            memory_usage,
            sram_efficiency,
            kernel_occupancy,
        };
        // Update global statistics
        {
            let mut global_stats = self.stats.write();
            global_stats.backward_time_us = backward_time_us;
        }
        info!("Flash Attention backward completed in {}μs", backward_time_us);
        Ok(FlashGradOutput {
            dq,
            dk,
            dv,
            stats,
        })
    }
    fn name(&self) -> &str {
        "FlashAttention"
    }
    fn supports_config(&self, config: &FlashAttentionConfig) -> bool {
        // Check if configuration is compatible with this backend
        config.validate().is_ok() &&
        config.num_heads <= 128 &&         // Reasonable upper limit
        config.head_dim <= 512 &&          // Reasonable upper limit
        config.max_seq_len <= 65_536        // 64K context limit
    }
    fn optimize_config(&self, mut config: FlashAttentionConfig) -> FlashResult<FlashAttentionConfig> {
        // Optimize block sizes based on GPU memory hierarchy
        let gpu_info = self.cuda_backend.get_device_info()
        .map_err(|e| FlashError::backend_init(format!("Failed to get GPU info: {}", e)))?;
        let shared_memory_size = gpu_info.shared_memory_per_block;
        let max_threads_per_block = gpu_info.max_threads_per_block;
        // Calculate optimal block sizes based on shared memory constraints
        let element_size = match config.precision {
            crate::config::PrecisionMode::FP32 => 4,
            crate::config::PrecisionMode::FP16 |
            crate::config::PrecisionMode::BF16 => 2,
            crate::config::PrecisionMode::FP8E4M3 { .. } |
            crate::config::PrecisionMode::FP8E5M2 { .. } => 1,
            crate::config::PrecisionMode::Mixed { storage_precision, .. } => {
                match storage_precision {
                    crate::config::Precision::FP32 => 4,
                    crate::config::Precision::FP16 |
                    crate::config::Precision::BF16 => 2,
                    crate::config::Precision::INT8 |
                    crate::config::Precision::FP8E4M3 |
                    crate::config::Precision::FP8E5M2 => 1,
                }
            }
        };
        // Calculate block sizes that fit in shared memory
        let memory_per_element = element_size * config.head_dim;
        let max_block_size = (shared_memory_size / (2 * memory_per_element)).min(256);
        // Round down to nearest multiple of 32 for coalescing
        let optimal_block_size = (max_block_size / 32) * 32;
        if optimal_block_size >= 32 {
            config.block_size_q = optimal_block_size;
            config.block_size_kv = optimal_block_size;
        }
        // Optimize memory pool based on estimated usage
        let estimated_memory = config.estimate_memory_usage(8, config.max_seq_len); // Estimate for batch size 8
        config.backend_config.cuda.memory_pool.initial_size = estimated_memory.max(
        config.backend_config.cuda.memory_pool.initial_size
        );
        info!("Optimized Flash Attention config: block_size={}x{}, memory_pool={}MB",
        config.block_size_q, config.block_size_kv,
        config.backend_config.cuda.memory_pool.initial_size / (1024 * 1024));
        Ok(config)
    }
 }
 /// Flash Attention factory for creating different backend implementations
 pub struct FlashAttentionFactory;
 impl FlashAttentionFactory {
    /// Create the best available Flash Attention backend for the given configuration
    pub fn create_backend(config: FlashAttentionConfig) -> FlashResult<Box<dyn FlashAttentionBackend>> {
        // For now, we only have the standard CUDA backend
        // In the future, this will select between CUDA, Quantum, Neuromorphic, and Edge backends
        let backend = FlashAttention::new(config)?;
        Ok(Box::new(backend))
    }
    /// Create a quantum-enhanced Flash Attention backend
    #[cfg(feature = "quantum")]
    pub fn create_quantum_backend(config: FlashAttentionConfig) -> FlashResult<Box<dyn FlashAttentionBackend>> {
        use crate::variants::quantum::QuantumFlashAttention;
        let backend = QuantumFlashAttention::new(config)?;
        Ok(Box::new(backend))
    }
    /// Create a neuromorphic Flash Attention backend
    #[cfg(feature = "neuromorphic")]
    pub fn create_neuromorphic_backend(config: FlashAttentionConfig) -> FlashResult<Box<dyn FlashAttentionBackend>> {
        use crate::variants::neuromorphic::NeuromorphicFlashAttention;
        let backend = NeuromorphicFlashAttention::new(config)?;
        Ok(Box::new(backend))
    }
    /// Create an edge-optimized Flash Attention backend
    #[cfg(feature = "edge")]
    pub fn create_edge_backend(config: FlashAttentionConfig) -> FlashResult<Box<dyn FlashAttentionBackend>> {
        use crate::variants::edge::EdgeFlashAttention;
        let backend = EdgeFlashAttention::new(config)?;
        Ok(Box::new(backend))
    }
    /// List all available backends
    pub fn available_backends() -> Vec<&'static str> {
        let mut backends = vec!["FlashAttention"];
        #[cfg(feature = "quantum")]
        backends.push("QuantumFlashAttention");
        #[cfg(feature = "neuromorphic")]
        backends.push("NeuromorphicFlashAttention");
        #[cfg(feature = "edge")]
        backends.push("EdgeFlashAttention");
        backends
    }
 }
 #[cfg(test)]
 mod tests {
    use super::*;
    use rtx_tensor::{Device, DType};
    #[tokio::test]
    async fn test_flash_attention_creation() {
        let config = FlashAttentionConfig::new(8, 64);
        // This will fail in CI without CUDA, but that's expected
        match FlashAttention::new(config) {
            Ok(flash) => {
                assert_eq!(flash.config().num_heads, 8);
                assert_eq!(flash.config().head_dim, 64);
            }
            Err(FlashError::BackendInit { .. }) => {
                // Expected in CI without CUDA
            }
            Err(e) => panic!("Unexpected error: {}", e),
        }
    }
    #[test]
    fn test_input_validation() {
        let config = FlashAttentionConfig::new(8, 64);
        // This will fail in CI, but we can test the error path
        if let Ok(flash) = FlashAttention::new(config) {
            // Create invalid tensors (different shapes)
            let q = Tensor::zeros(&[2, 8, 128, 64], DType::F32, Device::cuda(0).unwrap_or(Device::default())).unwrap();
            let k = Tensor::zeros(&[2, 8, 64, 64], DType::F32, Device::cuda(0).unwrap_or(Device::default())).unwrap(); // Wrong seq_len
            let v = Tensor::zeros(&[2, 8, 128, 64], DType::F32, Device::cuda(0).unwrap_or(Device::default())).unwrap();
            let result = flash.validate_inputs(&q, &k, &v);
            assert!(result.is_err());
        }
    }
    #[test]
    fn test_backend_factory() {
        let backends = FlashAttentionFactory::available_backends();
        assert!(backends.contains(&"FlashAttention"));
        #[cfg(feature = "quantum")]
        assert!(backends.contains(&"QuantumFlashAttention"));
        #[cfg(feature = "neuromorphic")]
        assert!(backends.contains(&"NeuromorphicFlashAttention"));
        #[cfg(feature = "edge")]
        assert!(backends.contains(&"EdgeFlashAttention"));
    }
 }
@@ -1,164 +0,0 @@
 //! # RTX Flash Attention
 //!
 //! Revolutionary Flash Attention implementation providing 5-8x speedup over existing solutions
 //! with O(n) memory complexity and support for quantum, neuromorphic, and edge computing variants.
 //!
 //! ## Features
 //!
 //! - **Core Flash Attention**: O(n) memory forward/backward passes with SRAM tiling
 //! - **Quantum-Enhanced**: Quantum pattern optimization for 2-3x additional speedup
 //! - **Neuromorphic**: Spike-based computation with 1000x energy efficiency
 //! - **Edge-Optimized**: Deployment to resource-constrained platforms
 //! - **32K+ Context**: Support for extremely long sequences
 //!
 //! ## Usage
 //!
 //! ```rust,no_run
 //! use rtx_flash_attention::{FlashAttention, FlashAttentionConfig};
 //! use rtx_tensor::Tensor;
 //!
 //! let config = FlashAttentionConfig {
    //!     num_heads: 32,
    //!     head_dim: 128,
    //!     block_size_q: 64,
    //!     block_size_kv: 64,
    //!     causal: true,
    //! };
    //!
    //! let flash_attention = FlashAttention::new(config)?;
    //!
    //! let (output, lse) = flash_attention.forward(&q, &k, &v)?;
    //! # Ok::<(), Box<dyn std::error::Error>>(())
    //! ```
    pub mod error;
    pub mod config;
    pub mod core;
    pub mod kernels;
    pub mod memory;
    #[cfg(any(feature = "quantum", feature = "neuromorphic", feature = "edge"))]
    pub mod variants {
        #[cfg(feature = "quantum")]
        pub mod quantum;
        #[cfg(feature = "neuromorphic")]
        pub mod neuromorphic;
        #[cfg(feature = "edge")]
        pub mod edge;
        #[cfg(feature = "quantum")]
        pub use quantum::*;
        #[cfg(feature = "neuromorphic")]
        pub use neuromorphic::*;
        #[cfg(feature = "edge")]
        pub use edge::*;
    }
    // Re-exports
    pub use error::{FlashError, FlashResult};
    pub use config::FlashAttentionConfig;
    pub use core::{FlashAttention, FlashAttentionBackend};
    pub use memory::{SRAMManager, BlockManager};
    use rtx_tensor::Tensor;
    use std::sync::Arc;
    use parking_lot::RwLock;
    /// Flash Attention execution statistics
    #[derive(Debug, Clone)]
    pub struct FlashStats {
        /// Forward pass execution time (microseconds)
        pub forward_time_us: u64,
        /// Backward pass execution time (microseconds)
        pub backward_time_us: u64,
        /// Memory usage (bytes)
        pub memory_usage: usize,
        /// SRAM efficiency (0.0 - 1.0)
        pub sram_efficiency: f32,
        /// Kernel occupancy (0.0 - 1.0)
        pub kernel_occupancy: f32,
    }
    /// Flash Attention output containing result and metadata
    #[derive(Debug)]
    pub struct FlashOutput {
        /// Attention output tensor [batch, heads, seq_len, head_dim]
        pub output: Tensor,
        /// Log-sum-exp for numerical stability [batch, heads, seq_len]
        pub lse: Tensor,
        /// Execution statistics
        pub stats: FlashStats,
    }
    /// Flash Attention gradient output
    #[derive(Debug)]
    pub struct FlashGradOutput {
        /// Query gradients [batch, heads, seq_len, head_dim]
        pub dq: Tensor,
        /// Key gradients [batch, heads, seq_len, head_dim]
        pub dk: Tensor,
        /// Value gradients [batch, heads, seq_len, head_dim]
        pub dv: Tensor,
        /// Execution statistics
        pub stats: FlashStats,
    }
    /// Global Flash Attention registry for backend management
    pub struct FlashRegistry {
        backends: Arc<RwLock<Vec<Box<dyn FlashAttentionBackend + Send + Sync>>>>,
    }
    impl FlashRegistry {
        /// Create a new Flash Attention registry
        pub fn new() -> Self {
            Self {
                backends: Arc::new(RwLock::new(Vec::new())),
            }
        }
        /// Register a new Flash Attention backend
        pub fn register_backend<B>(&self, backend: B)
        where
        B: FlashAttentionBackend + Send + Sync + 'static,
        {
            self.backends.write().push(Box::new(backend));
        }
        /// Get the best available backend for the given configuration
        pub fn get_backend(&self, config: &FlashAttentionConfig) -> Option<&dyn FlashAttentionBackend> {
            // For now, return the first backend. In the future, we could implement
            // sophisticated backend selection based on hardware capabilities and workload
            self.backends.read().first().map(|b| b.as_ref())
        }
    }
    impl Default for FlashRegistry {
        fn default() -> Self {
            Self::new()
        }
    }
    /// Global Flash Attention registry instance
    static FLASH_REGISTRY: std::sync::OnceLock<FlashRegistry> = std::sync::OnceLock::new();
    /// Get the global Flash Attention registry
    pub fn global_registry() -> &'static FlashRegistry {
        FLASH_REGISTRY.get_or_init(FlashRegistry::default)
    }
    #[cfg(test)]
    mod tests {
        use super::*;
        #[test]
        fn test_flash_registry() {
            let registry = FlashRegistry::new();
            assert!(registry.backends.read().is_empty());
        }
        #[test]
        fn test_global_registry() {
            let registry = global_registry();
            assert!(registry.backends.read().is_empty());
        }
    }
@@ -1,67 +0,0 @@
 //! Curriculum-aware data loader.
 use super::sample::Sample;
 use super::scorers::DifficultyScorer;
 use super::strategies::CurriculumStrategy;
 use super::schedules::Schedule;
 use super::state::CurriculumState;
 /// Data loader with curriculum learning integration
 pub struct CurriculumDataLoader<S: Sample> {
    samples: Vec<S>,
    strategy: Box<dyn CurriculumStrategy>,
    difficulty_scorer: Box<dyn DifficultyScorer>,
    schedule: Box<dyn Schedule>,
    batch_size: usize,
    curriculum_state: CurriculumState,
 }
 impl<S: Sample> CurriculumDataLoader<S> {
    pub fn new(
        samples: Vec<S>,
        strategy: Box<dyn CurriculumStrategy>,
        difficulty_scorer: Box<dyn DifficultyScorer>,
        schedule: Box<dyn Schedule>,
        batch_size: usize,
    ) -> Self {
        Self {
            samples,
            strategy,
            difficulty_scorer,
            schedule,
            batch_size,
            curriculum_state: CurriculumState::new(),
        }
    }
    pub fn next_batch(&mut self) -> Option<Vec<S>> {
        if self.samples.is_empty() {
            return None;
        }
        // Update difficulty threshold from schedule
        self.curriculum_state.difficulty_threshold =
            self.schedule.get_difficulty_at_step(self.curriculum_state.step);
        // Update strategy's threshold
        self.strategy.update_difficulty_threshold(&mut self.curriculum_state);
        // Select samples using strategy
        let batch = self.strategy.select_samples(
            &self.samples,
            self.difficulty_scorer.as_ref(),
            &mut self.curriculum_state,
            self.batch_size,
        );
        Some(batch)
    }
    pub fn step(&mut self) {
        self.curriculum_state.step += 1;
    }
    pub fn get_curriculum_state(&self) -> &CurriculumState {
        &self.curriculum_state
    }
 }
@@ -1,18 +0,0 @@
 //! Curriculum learning errors.
 use thiserror::Error;
 /// Errors that can occur during curriculum learning
 #[derive(Error, Debug)]
 pub enum CurriculumError {
    #[error("Invalid difficulty score: {0}")]
    InvalidDifficultyScore(f32),
    #[error("Invalid schedule parameters: {0}")]
    InvalidScheduleParameters(String),
    #[error("Sample not found: {0}")]
    SampleNotFound(usize),
    #[error("Strategy configuration error: {0}")]
    StrategyConfiguration(String),
    #[error("Performance tracking error: {0}")]
    PerformanceTracking(String),
 }
@@ -1,86 +0,0 @@
 //! Curriculum learning metrics.
 /// Comprehensive metrics for curriculum learning
 #[derive(Debug, Clone)]
 pub struct CurriculumMetrics {
    pub batches_processed: usize,
    pub total_samples_seen: usize,
    pub samples_selected: usize,
    pub difficulty_distributions: Vec<Vec<f32>>,
    pub performance_history: Vec<f32>,
    pub selection_history: Vec<usize>,
 }
 impl CurriculumMetrics {
    pub fn new() -> Self {
        Self {
            batches_processed: 0,
            total_samples_seen: 0,
            samples_selected: 0,
            difficulty_distributions: Vec::new(),
            performance_history: Vec::new(),
            selection_history: Vec::new(),
        }
    }
    pub fn record_batch_difficulty_distribution(&mut self, difficulties: Vec<f32>) {
        self.batches_processed += 1;
        self.total_samples_seen += difficulties.len();
        self.difficulty_distributions.push(difficulties);
    }
    pub fn record_sample_selection_stats(&mut self, total_samples: usize, selected_samples: usize) {
        self.total_samples_seen += total_samples;
        self.samples_selected += selected_samples;
        self.selection_history.push(selected_samples);
    }
    pub fn update_performance_tracking(&mut self, performance: f32) {
        self.performance_history.push(performance);
    }
    pub fn get_selection_stats(&self) -> SelectionStats {
        let selection_ratio = if self.total_samples_seen > 0 {
            self.samples_selected as f32 / self.total_samples_seen as f32
        } else {
            0.0
        };
        let avg_difficulty = if !self.difficulty_distributions.is_empty() {
            let total_difficulty: f32 = self.difficulty_distributions
                .iter()
                .flat_map(|d| d.iter())
                .sum();
            let total_count: usize = self.difficulty_distributions
                .iter()
                .map(|d| d.len())
                .sum();
            total_difficulty / total_count as f32
        } else {
            0.0
        };
        SelectionStats {
            selection_ratio,
            avg_difficulty,
            batches_processed: self.batches_processed,
        }
    }
    pub fn reset(&mut self) {
        *self = Self::new();
    }
 }
 impl Default for CurriculumMetrics {
    fn default() -> Self {
        Self::new()
    }
 }
 #[derive(Debug, Clone)]
 pub struct SelectionStats {
    pub selection_ratio: f32,
    pub avg_difficulty: f32,
    pub batches_processed: usize,
 }
@@ -1,67 +0,0 @@
 //! Curriculum optimizer for hyperparameter tuning.
 use std::collections::HashMap;
 /// Curriculum optimizer for hyperparameter tuning
 #[derive(Debug)]
 pub struct CurriculumOptimizer {
    configuration_performance: HashMap<CurriculumConfig, f32>,
 }
 type OrderedFloat = ordered_float::OrderedFloat<f32>;
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct CurriculumConfig {
    pub strategy_type: CurriculumStrategyType,
    pub schedule_type: ScheduleType,
    pub parameters: Vec<OrderedFloat>,
 }
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum CurriculumStrategyType {
    EasyToHard,
    AntiCurriculum,
    SelfPaced,
    CompetencyBased,
    DataDriven,
 }
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ScheduleType {
    Linear,
    Exponential,
    Adaptive,
    Cyclic,
    MultiTask,
 }
 impl CurriculumOptimizer {
    pub fn new() -> Self {
        Self {
            configuration_performance: HashMap::new(),
        }
    }
    pub fn record_configuration_performance(&mut self, config: CurriculumConfig, performance: f32) {
        self.configuration_performance.insert(config, performance);
    }
    pub fn get_best_configuration(&self) -> Option<CurriculumConfig> {
        self.configuration_performance
            .iter()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .map(|(config, _)| config.clone())
    }
    pub fn optimize_hyperparameters(&self, base_config: &CurriculumConfig, _iterations: usize) -> Option<CurriculumConfig> {
        // Simple optimization: return best known configuration
        // In a real implementation, this would use more sophisticated optimization
        self.get_best_configuration().or_else(|| Some(base_config.clone()))
    }
 }
 impl Default for CurriculumOptimizer {
    fn default() -> Self {
        Self::new()
    }
 }
@@ -1,12 +0,0 @@
 //! Sample trait for curriculum learning.
 use std::collections::HashMap;
 /// Trait for training samples that can be used with curriculum learning
 pub trait Sample: Clone {
    /// Unique identifier for the sample
    fn id(&self) -> usize;
    /// Extract complexity features for difficulty scoring
    fn complexity_features(&self) -> HashMap<String, f32>;
 }
@@ -1,119 +0,0 @@
 //! Curriculum-aware batch sampler.
 use std::collections::HashMap;
 use super::sample::Sample;
 /// Curriculum-aware batch sampler
 pub struct CurriculumBatchSampler {
    batch_size: usize,
    temperature: f32,
    sampling_strategy: SamplingStrategy,
 }
 #[derive(Debug, Clone)]
 pub enum SamplingStrategy {
    WeightedRandom,
    Uniform,
    ThresholdBased,
 }
 impl CurriculumBatchSampler {
    pub fn new(batch_size: usize, temperature: f32, sampling_strategy: SamplingStrategy) -> Self {
        Self {
            batch_size,
            temperature,
            sampling_strategy,
        }
    }
    pub fn sample_batch<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scores: &HashMap<usize, f32>,
        difficulty_threshold: f32,
    ) -> Vec<usize> {
        match self.sampling_strategy {
            SamplingStrategy::WeightedRandom => {
                self.weighted_random_sampling(samples, difficulty_scores, difficulty_threshold)
            }
            SamplingStrategy::Uniform => {
                self.uniform_sampling(samples, difficulty_scores, difficulty_threshold)
            }
            SamplingStrategy::ThresholdBased => {
                self.threshold_based_sampling(samples, difficulty_scores, difficulty_threshold)
            }
        }
    }
    fn weighted_random_sampling<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scores: &HashMap<usize, f32>,
        _difficulty_threshold: f32,
    ) -> Vec<usize> {
        // Simple implementation: select based on inverse difficulty weights
        let mut indices = Vec::new();
        let total_samples = self.batch_size.min(samples.len());
        for (idx, sample) in samples.iter().enumerate() {
            if indices.len() >= total_samples {
                break;
            }
            let score = difficulty_scores.get(&sample.id()).copied().unwrap_or(0.5);
            let weight = 1.0 / (1.0 + score * self.temperature);
            // Simple selection based on weight (in real implementation would use proper random sampling)
            if weight > 0.5 {
                indices.push(idx);
            }
        }
        // Fill up to batch_size if needed
        while indices.len() < total_samples && indices.len() < samples.len() {
            for (idx, _) in samples.iter().enumerate() {
                if !indices.contains(&idx) {
                    indices.push(idx);
                    if indices.len() >= total_samples {
                        break;
                    }
                }
            }
            break;
        }
        indices
    }
    fn uniform_sampling<S: Sample>(
        &self,
        samples: &[S],
        _difficulty_scores: &HashMap<usize, f32>,
        _difficulty_threshold: f32,
    ) -> Vec<usize> {
        let total_samples = self.batch_size.min(samples.len());
        (0..total_samples).collect()
    }
    fn threshold_based_sampling<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scores: &HashMap<usize, f32>,
        difficulty_threshold: f32,
    ) -> Vec<usize> {
        let mut indices = Vec::new();
        for (idx, sample) in samples.iter().enumerate() {
            if indices.len() >= self.batch_size {
                break;
            }
            let score = difficulty_scores.get(&sample.id()).copied().unwrap_or(0.5);
            if score <= difficulty_threshold {
                indices.push(idx);
            }
        }
        indices
    }
 }
@@ -1,208 +0,0 @@
 //! Curriculum schedule implementations.
 use std::collections::HashMap;
 /// Trait for curriculum scheduling
 pub trait Schedule: Send + Sync {
    /// Get difficulty threshold at given training step
    fn get_difficulty_at_step(&self, step: usize) -> f32;
 }
 /// Linear curriculum schedule
 #[derive(Debug, Clone)]
 pub struct LinearSchedule {
    initial_difficulty: f32,
    final_difficulty: f32,
    total_steps: usize,
 }
 impl LinearSchedule {
    pub fn new(initial_difficulty: f32, final_difficulty: f32, total_steps: usize) -> Self {
        Self {
            initial_difficulty,
            final_difficulty,
            total_steps,
        }
    }
 }
 impl Schedule for LinearSchedule {
    fn get_difficulty_at_step(&self, step: usize) -> f32 {
        if step >= self.total_steps {
            return self.final_difficulty;
        }
        let progress = step as f32 / self.total_steps as f32;
        self.initial_difficulty + progress * (self.final_difficulty - self.initial_difficulty)
    }
 }
 /// Exponential curriculum schedule
 #[derive(Debug, Clone)]
 pub struct ExponentialSchedule {
    initial_difficulty: f32,
    final_difficulty: f32,
    growth_rate: f32,
 }
 impl ExponentialSchedule {
    pub fn new(initial_difficulty: f32, final_difficulty: f32, growth_rate: f32) -> Self {
        Self {
            initial_difficulty,
            final_difficulty,
            growth_rate,
        }
    }
 }
 impl Schedule for ExponentialSchedule {
    fn get_difficulty_at_step(&self, step: usize) -> f32 {
        let exp_factor = 1.0 - (-self.growth_rate * step as f32).exp();
        let difficulty = self.initial_difficulty + exp_factor * (self.final_difficulty - self.initial_difficulty);
        difficulty.min(self.final_difficulty)
    }
 }
 /// Adaptive curriculum schedule (adjusts based on performance)
 #[derive(Debug, Clone)]
 pub struct AdaptiveSchedule {
    initial_difficulty: f32,
    final_difficulty: f32,
    adaptation_rate: f32,
    target_performance: f32,
    performance_window: Vec<f32>,
    current_difficulty: f32,
 }
 impl AdaptiveSchedule {
    pub fn new(
        initial_difficulty: f32,
        final_difficulty: f32,
        adaptation_rate: f32,
        target_performance: f32,
    ) -> Self {
        Self {
            initial_difficulty,
            final_difficulty,
            adaptation_rate,
            target_performance,
            performance_window: Vec::new(),
            current_difficulty: initial_difficulty,
        }
    }
    pub fn update_performance(&mut self, performance: f32) {
        self.performance_window.push(performance);
        if self.performance_window.len() > 10 {
            self.performance_window.remove(0);
        }
    }
    fn current_performance(&self) -> f32 {
        if self.performance_window.is_empty() {
            self.target_performance
        } else {
            self.performance_window.iter().sum::<f32>() / self.performance_window.len() as f32
        }
    }
 }
 impl Schedule for AdaptiveSchedule {
    fn get_difficulty_at_step(&self, _step: usize) -> f32 {
        let current_perf = self.current_performance();
        let adjustment = if current_perf > self.target_performance {
            self.adaptation_rate // Increase difficulty
        } else {
            -self.adaptation_rate // Decrease difficulty
        };
        (self.current_difficulty + adjustment).clamp(self.initial_difficulty, self.final_difficulty)
    }
 }
 /// Cyclic curriculum schedule
 #[derive(Debug, Clone)]
 pub struct CyclicSchedule {
    min_difficulty: f32,
    max_difficulty: f32,
    cycle_length: usize,
 }
 impl CyclicSchedule {
    pub fn new(min_difficulty: f32, max_difficulty: f32, cycle_length: usize) -> Self {
        Self {
            min_difficulty,
            max_difficulty,
            cycle_length,
        }
    }
 }
 impl Schedule for CyclicSchedule {
    fn get_difficulty_at_step(&self, step: usize) -> f32 {
        let cycle_position = (step % self.cycle_length) as f32;
        let normalized_position = cycle_position / self.cycle_length as f32;
        // Triangle wave: goes from min to max and back
        let triangle_wave = if normalized_position <= 0.5 {
            normalized_position * 2.0 // 0 -> 1
        } else {
            2.0 - normalized_position * 2.0 // 1 -> 0
        };
        self.min_difficulty + triangle_wave * (self.max_difficulty - self.min_difficulty)
    }
 }
 /// Multi-task curriculum schedule coordinator
 #[derive(Debug)]
 pub struct MultiTaskSchedule {
    task_schedules: HashMap<String, Box<dyn Schedule>>,
    coordination_weight: f32,
 }
 impl MultiTaskSchedule {
    pub fn new() -> Self {
        Self {
            task_schedules: HashMap::new(),
            coordination_weight: 0.0,
        }
    }
    pub fn add_task(&mut self, task_name: &str, schedule: Box<dyn Schedule>) {
        self.task_schedules.insert(task_name.to_string(), schedule);
    }
    pub fn set_coordination_weight(&mut self, weight: f32) {
        self.coordination_weight = weight;
    }
    pub fn get_task_difficulty(&self, task_name: &str, step: usize) -> Option<f32> {
        self.task_schedules.get(task_name).map(|schedule| schedule.get_difficulty_at_step(step))
    }
    pub fn get_coordinated_difficulty(&self, task_name: &str, step: usize) -> Option<f32> {
        let task_difficulty = self.get_task_difficulty(task_name, step)?;
        if self.coordination_weight <= 0.0 {
            return Some(task_difficulty);
        }
        // Calculate average difficulty across all tasks
        let total_difficulty: f32 = self.task_schedules
            .values()
            .map(|schedule| schedule.get_difficulty_at_step(step))
            .sum();
        let avg_difficulty = total_difficulty / self.task_schedules.len() as f32;
        // Blend task-specific and average difficulty
        Some(task_difficulty * (1.0 - self.coordination_weight) + avg_difficulty * self.coordination_weight)
    }
 }
 impl Default for MultiTaskSchedule {
    fn default() -> Self {
        Self::new()
    }
 }
@@ -1,182 +0,0 @@
 //! Difficulty scoring implementations.
 use std::collections::HashMap;
 use super::sample::Sample;
 /// Trait for difficulty scoring functions (object-safe version)
 pub trait DifficultyScorer: Send + Sync {
    /// Score a sample's difficulty (0.0 = easy, 1.0 = hard) using complexity features
    fn score_features(&self, features: &HashMap<String, f32>) -> f32;
    /// Optional: Update scorer based on performance feedback
    fn update_from_performance(&mut self, _sample_id: usize, _performance: f32) {}
 }
 /// Convenience method for scoring samples directly
 pub trait DifficultyExt<S: Sample> {
    fn score(&self, sample: &S) -> f32;
 }
 impl<T: DifficultyScorer, S: Sample> DifficultyExt<S> for T {
    fn score(&self, sample: &S) -> f32 {
        let features = sample.complexity_features();
        self.score_features(&features)
    }
 }
 /// Length-based difficulty scorer - longer sequences are harder
 #[derive(Debug, Clone)]
 pub struct LengthBasedDifficultyScorer {
    max_length: f32,
 }
 impl LengthBasedDifficultyScorer {
    pub fn new() -> Self {
        Self { max_length: 100.0 }
    }
    pub fn with_max_length(max_length: f32) -> Self {
        Self { max_length }
    }
 }
 impl Default for LengthBasedDifficultyScorer {
    fn default() -> Self {
        Self::new()
    }
 }
 impl DifficultyScorer for LengthBasedDifficultyScorer {
    fn score_features(&self, features: &HashMap<String, f32>) -> f32 {
        let length = features.get("sequence_length").copied().unwrap_or(0.0);
        (length / self.max_length).min(1.0)
    }
 }
 /// Variance-based difficulty scorer - higher variance indicates more complexity
 #[derive(Debug, Clone)]
 pub struct VarianceBasedDifficultyScorer {
    max_variance: f32,
 }
 impl VarianceBasedDifficultyScorer {
    pub fn new() -> Self {
        Self { max_variance: 10.0 }
    }
    pub fn with_max_variance(max_variance: f32) -> Self {
        Self { max_variance }
    }
 }
 impl Default for VarianceBasedDifficultyScorer {
    fn default() -> Self {
        Self::new()
    }
 }
 impl DifficultyScorer for VarianceBasedDifficultyScorer {
    fn score_features(&self, features: &HashMap<String, f32>) -> f32 {
        let variance = features.get("variance").copied().unwrap_or(0.0);
        (variance / self.max_variance).min(1.0)
    }
 }
 /// Composite difficulty scorer that combines multiple scorers
 #[derive(Debug)]
 pub struct CompositeDifficultyScorer {
    scorers: Vec<Box<dyn DifficultyScorer>>,
    weights: Vec<f32>,
 }
 impl CompositeDifficultyScorer {
    pub fn new() -> Self {
        Self {
            scorers: Vec::new(),
            weights: Vec::new(),
        }
    }
    pub fn add_scorer(&mut self, scorer: Box<dyn DifficultyScorer>, weight: f32) {
        self.scorers.push(scorer);
        self.weights.push(weight);
        self.normalize_weights();
    }
    pub fn get_weights(&self) -> &[f32] {
        &self.weights
    }
    fn normalize_weights(&mut self) {
        let sum: f32 = self.weights.iter().sum();
        if sum > 0.0 {
            for weight in &mut self.weights {
                *weight /= sum;
            }
        }
    }
 }
 impl Default for CompositeDifficultyScorer {
    fn default() -> Self {
        Self::new()
    }
 }
 impl DifficultyScorer for CompositeDifficultyScorer {
    fn score_features(&self, features: &HashMap<String, f32>) -> f32 {
        if self.scorers.is_empty() {
            return 0.5; // Default difficulty
        }
        self.scorers
            .iter()
            .zip(&self.weights)
            .map(|(scorer, &weight)| scorer.score_features(features) * weight)
            .sum()
    }
 }
 /// Curriculum-based difficulty scorer using performance history
 #[derive(Debug, Clone)]
 pub struct CurriculumBasedDifficultyScorer {
    performance_history: HashMap<usize, f32>,
    default_difficulty: f32,
 }
 impl CurriculumBasedDifficultyScorer {
    pub fn new() -> Self {
        Self {
            performance_history: HashMap::new(),
            default_difficulty: 0.5,
        }
    }
    pub fn update_performance_history(&mut self, sample_id: usize, performance: f32) {
        self.performance_history.insert(sample_id, performance);
    }
 }
 impl Default for CurriculumBasedDifficultyScorer {
    fn default() -> Self {
        Self::new()
    }
 }
 impl DifficultyScorer for CurriculumBasedDifficultyScorer {
    fn score_features(&self, features: &HashMap<String, f32>) -> f32 {
        // For curriculum-based scoring, we need the sample ID from features
        if let Some(&sample_id) = features.get("sample_id") {
            self.performance_history
                .get(&(sample_id as usize))
                .map(|&perf| 1.0 - perf) // Inverse relationship: low performance = high difficulty
                .unwrap_or(self.default_difficulty)
        } else {
            self.default_difficulty
        }
    }
    fn update_from_performance(&mut self, sample_id: usize, performance: f32) {
        self.performance_history.insert(sample_id, performance);
    }
 }
@@ -1,336 +0,0 @@
 //! Simple curriculum learning test without external dependencies
 use std::collections::HashMap;
 /// Test sample implementation
 #[derive(Debug, Clone, PartialEq)]
 pub struct TestSample {
    pub id: usize,
    pub data: Vec<f32>,
    pub metadata: HashMap<String, f32>,
 }
 impl TestSample {
    pub fn new(id: usize, data: Vec<f32>) -> Self {
        Self {
            id,
            data,
            metadata: HashMap::new(),
        }
    }
    pub fn with_metadata(mut self, key: &str, value: f32) -> Self {
        self.metadata.insert(key.to_string(), value);
        self
    }
    pub fn id(&self) -> usize {
        self.id
    }
    pub fn complexity_features(&self) -> HashMap<String, f32> {
        let mut features = self.metadata.clone();
        features.insert("sequence_length".to_string(), self.data.len() as f32);
        features.insert("variance".to_string(), 
            self.data.iter().map(|&x| x * x).sum::<f32>() / self.data.len() as f32);
        features
    }
 }
 /// Length-based difficulty scorer
 #[derive(Debug, Clone)]
 pub struct LengthBasedDifficultyScorer {
    max_length: f32,
 }
 impl LengthBasedDifficultyScorer {
    pub fn new() -> Self {
        Self { max_length: 100.0 }
    }
    pub fn score(&self, sample: &TestSample) -> f32 {
        let features = sample.complexity_features();
        let length = features.get("sequence_length").copied().unwrap_or(0.0);
        (length / self.max_length).min(1.0)
    }
 }
 /// Curriculum state
 #[derive(Debug, Clone)]
 pub struct CurriculumState {
    pub step: usize,
    pub difficulty_threshold: f32,
 }
 impl CurriculumState {
    pub fn new() -> Self {
        Self {
            step: 0,
            difficulty_threshold: 0.0,
        }
    }
 }
 /// Easy-to-hard strategy
 #[derive(Debug, Clone)]
 pub struct EasyToHardStrategy {
    initial_threshold: f32,
    threshold_increment: f32,
 }
 impl EasyToHardStrategy {
    pub fn new(initial_threshold: f32, threshold_increment: f32) -> Self {
        Self {
            initial_threshold,
            threshold_increment,
        }
    }
    pub fn select_samples(
        &self,
        samples: &[TestSample],
        difficulty_scorer: &LengthBasedDifficultyScorer,
        curriculum_state: &CurriculumState,
        batch_size: usize,
    ) -> Vec<TestSample> {
        let mut scored_samples: Vec<(TestSample, f32)> = samples
            .iter()
            .map(|s| (s.clone(), difficulty_scorer.score(s)))
            .collect();
        // Filter by current difficulty threshold
        scored_samples.retain(|(_, score)| *score <= curriculum_state.difficulty_threshold);
        // Sort by difficulty (easiest first)
        scored_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
        // If not enough samples, include some harder ones
        if scored_samples.len() < batch_size {
            let mut all_samples: Vec<(TestSample, f32)> = samples
                .iter()
                .map(|s| (s.clone(), difficulty_scorer.score(s)))
                .collect();
            all_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
            scored_samples = all_samples.into_iter().take(batch_size).collect();
        }
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    pub fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        curriculum_state.difficulty_threshold = 
            (self.initial_threshold + curriculum_state.step as f32 * self.threshold_increment).min(1.0);
    }
 }
 /// Linear schedule
 #[derive(Debug, Clone)]
 pub struct LinearSchedule {
    initial_difficulty: f32,
    final_difficulty: f32,
    total_steps: usize,
 }
 impl LinearSchedule {
    pub fn new(initial_difficulty: f32, final_difficulty: f32, total_steps: usize) -> Self {
        Self {
            initial_difficulty,
            final_difficulty,
            total_steps,
        }
    }
    pub fn get_difficulty_at_step(&self, step: usize) -> f32 {
        if step >= self.total_steps {
            return self.final_difficulty;
        }
        let progress = step as f32 / self.total_steps as f32;
        self.initial_difficulty + progress * (self.final_difficulty - self.initial_difficulty)
    }
 }
 fn main() {
    println!("🎯 Testing RTX Curriculum Learning Implementation (Simple)");
    println!("=========================================================");
    // Test 1: Difficulty Scoring
    println!("\n📊 Test 1: Difficulty Scoring");
    let scorer = LengthBasedDifficultyScorer::new();
    let easy_sample = TestSample::new(1, vec![1.0, 2.0]);
    let hard_sample = TestSample::new(2, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
    let easy_score = scorer.score(&easy_sample);
    let hard_score = scorer.score(&hard_sample);
    println!("   Easy sample (len=2): difficulty = {:.3}", easy_score);
    println!("   Hard sample (len=5): difficulty = {:.3}", hard_score);
    assert!(easy_score < hard_score, "❌ Easy sample should have lower difficulty");
    assert!(easy_score >= 0.0 && easy_score <= 1.0, "❌ Easy score should be normalized");
    assert!(hard_score >= 0.0 && hard_score <= 1.0, "❌ Hard score should be normalized");
    println!("   ✅ Difficulty scoring works correctly");
    // Test 2: Curriculum Strategy
    println!("\n📈 Test 2: Easy-to-Hard Curriculum Strategy");
    let strategy = EasyToHardStrategy::new(0.1, 0.05);
    let samples = vec![
        TestSample::new(1, vec![1.0; 10]), // Hard
        TestSample::new(2, vec![1.0; 2]),  // Easy
        TestSample::new(3, vec![1.0; 5]),  // Medium
    ];
    let mut state = CurriculumState::new();
    // Initially should select mostly easy samples
    let selected = strategy.select_samples(&samples, &scorer, &state, 2);
    println!("   Selected {} samples initially", selected.len());
    let selected_lengths: Vec<usize> = selected.iter()
        .map(|s| s.data.len())
        .collect();
    println!("   Selected sample lengths: {:?}", selected_lengths);
    assert_eq!(selected.len(), 2, "❌ Should select requested batch size");
    assert!(selected_lengths.contains(&2), "❌ Easy sample should be selected initially");
    println!("   ✅ Initial selection prefers easy samples");
    // Advance curriculum
    for i in 0..20 {
        state.step += 1;
        strategy.update_difficulty_threshold(&mut state);
        println!("   Step {}: threshold = {:.3}", i + 1, state.difficulty_threshold);
    }
    let selected_later = strategy.select_samples(&samples, &scorer, &state, 2);
    println!("   Selected {} samples after progression", selected_later.len());
    assert_eq!(selected_later.len(), 2, "❌ Should still select requested batch size");
    println!("   ✅ Curriculum progression works");
    // Test 3: Linear Schedule
    println!("\n⏱️  Test 3: Linear Schedule");
    let schedule = LinearSchedule::new(0.1, 0.9, 100);
    let step_0 = schedule.get_difficulty_at_step(0);
    let step_50 = schedule.get_difficulty_at_step(50);
    let step_100 = schedule.get_difficulty_at_step(100);
    let step_150 = schedule.get_difficulty_at_step(150);
    println!("   Step 0: difficulty = {:.3}", step_0);
    println!("   Step 50: difficulty = {:.3}", step_50);
    println!("   Step 100: difficulty = {:.3}", step_100);
    println!("   Step 150: difficulty = {:.3}", step_150);
    let epsilon = 1e-6;
    assert!((step_0 - 0.1).abs() < epsilon, "❌ Step 0 should be initial difficulty");
    assert!((step_50 - 0.5).abs() < epsilon, "❌ Step 50 should be halfway");
    assert!((step_100 - 0.9).abs() < epsilon, "❌ Step 100 should be final difficulty");
    assert!((step_150 - 0.9).abs() < epsilon, "❌ Step 150 should clamp at final difficulty");
    println!("   ✅ Linear schedule works correctly");
    // Test 4: Integration Test
    println!("\n🔗 Test 4: Integration Test");
    println!("   Testing complete curriculum learning pipeline...");
    let curriculum_samples = vec![
        TestSample::new(1, vec![1.0; 1]),   // Very easy
        TestSample::new(2, vec![1.0; 3]),   // Easy
        TestSample::new(3, vec![1.0; 5]),   // Medium
        TestSample::new(4, vec![1.0; 8]),   // Hard
        TestSample::new(5, vec![1.0; 12]),  // Very hard
    ];
    let mut curriculum_state = CurriculumState::new();
    let curriculum_strategy = EasyToHardStrategy::new(0.1, 0.1);
    let curriculum_schedule = LinearSchedule::new(0.1, 0.8, 10);
    for step in 0..10 {
        curriculum_state.step = step;
        curriculum_state.difficulty_threshold = curriculum_schedule.get_difficulty_at_step(step);
        curriculum_strategy.update_difficulty_threshold(&mut curriculum_state);
        let batch = curriculum_strategy.select_samples(
            &curriculum_samples,
            &scorer,
            &curriculum_state,
            2
        );
        let avg_length: f32 = batch.iter()
            .map(|s| s.data.len() as f32)
            .sum::<f32>() / batch.len() as f32;
        println!("   Step {}: threshold = {:.2}, avg_batch_length = {:.1}", 
                step, curriculum_state.difficulty_threshold, avg_length);
    }
    println!("   ✅ Complete pipeline integration successful");
    println!("\n🎉 All Curriculum Learning Tests Passed!");
    println!("   ✅ Difficulty scoring functions work correctly");
    println!("   ✅ Easy-to-hard strategy implemented properly");
    println!("   ✅ Linear scheduling functions correctly");
    println!("   ✅ Complete pipeline integration successful");
    println!("\n📋 Implementation Status:");
    println!("   🟢 Core curriculum learning framework - COMPLETE");
    println!("   🟢 Difficulty scoring strategies - COMPLETE");
    println!("   🟢 Curriculum selection strategies - COMPLETE");
    println!("   🟢 Scheduling algorithms - COMPLETE");
    println!("   🟢 State management - COMPLETE");
    println!("\n✨ RTX Curriculum Learning implementation follows strict TDD:");
    println!("   1. ✅ Tests written first (RED phase)");
    println!("   2. ✅ Minimal implementation created (GREEN phase)");
    println!("   3. ✅ Ready for refactoring phase");
    println!("   4. ✅ No mocks, stubs, or TODOs - production ready!");
 }
 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
    use super::*;
    #[test]
    fn test_difficulty_scoring() {
        let scorer = LengthBasedDifficultyScorer::new();
        let easy_sample = TestSample::new(1, vec![1.0, 2.0]);
        let hard_sample = TestSample::new(2, vec![1.0, 2.0, 3.0, 4.0, 5.0]);
        let easy_score = scorer.score(&easy_sample);
        let hard_score = scorer.score(&hard_sample);
        assert!(easy_score < hard_score);
        assert!(easy_score >= 0.0 && easy_score <= 1.0);
        assert!(hard_score >= 0.0 && hard_score <= 1.0);
    }
    #[test]
    fn test_curriculum_strategy() {
        let strategy = EasyToHardStrategy::new(0.1, 0.05);
        let scorer = LengthBasedDifficultyScorer::new();
        let samples = vec![
            TestSample::new(1, vec![1.0; 10]), // Hard
            TestSample::new(2, vec![1.0; 2]),  // Easy
            TestSample::new(3, vec![1.0; 5]),  // Medium
        ];
        let state = CurriculumState::new();
        let selected = strategy.select_samples(&samples, &scorer, &state, 2);
        assert_eq!(selected.len(), 2);
    }
    #[test]
    fn test_linear_schedule() {
        let schedule = LinearSchedule::new(0.1, 0.9, 100);
        assert!((schedule.get_difficulty_at_step(0) - 0.1).abs() < 1e-6);
        assert!((schedule.get_difficulty_at_step(50) - 0.5).abs() < 1e-6);
        assert!((schedule.get_difficulty_at_step(100) - 0.9).abs() < 1e-6);
    }
 }
@@ -1,385 +0,0 @@
 //! Standalone test for curriculum learning implementation
 //! This test runs independently without depending on the full RTX ecosystem
 use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 // Mock tensor types for testing
 #[derive(Debug, Clone)]
 pub struct MockDevice;
 #[derive(Debug, Clone)]
 pub struct MockTensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
 }
 impl MockTensor {
    pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
        Self { data, shape }
    }
 }
 /// Minimal error type for testing
 #[derive(thiserror::Error, Debug)]
 pub enum MockError {
    #[error("Test error: {0}")]
    TestError(String),
 }
 type Result<T> = std::result::Result<T, MockError>;
 /// Sample trait implementation
 pub trait Sample: Clone {
    fn id(&self) -> usize;
    fn complexity_features(&self) -> HashMap<String, f32>;
 }
 /// Test sample implementation
 #[derive(Debug, Clone, PartialEq)]
 pub struct TestSample {
    pub id: usize,
    pub data: Vec<f32>,
    pub label: Option<usize>,
    pub metadata: HashMap<String, f32>,
 }
 impl TestSample {
    pub fn new(id: usize, data: Vec<f32>, label: Option<usize>) -> Self {
        Self {
            id,
            data,
            label,
            metadata: HashMap::new(),
        }
    }
    pub fn with_metadata(mut self, key: &str, value: f32) -> Self {
        self.metadata.insert(key.to_string(), value);
        self
    }
 }
 impl Sample for TestSample {
    fn id(&self) -> usize {
        self.id
    }
    fn complexity_features(&self) -> HashMap<String, f32> {
        let mut features = self.metadata.clone();
        features.insert("sequence_length".to_string(), self.data.len() as f32);
        features.insert("variance".to_string(), 
            self.data.iter().map(|&x| x * x).sum::<f32>() / self.data.len() as f32);
        features
    }
 }
 /// Difficulty scorer trait
 pub trait DifficultyScorer: Send + Sync {
    fn score<S: Sample>(&self, sample: &S) -> f32;
 }
 /// Length-based difficulty scorer
 #[derive(Debug, Clone)]
 pub struct LengthBasedDifficultyScorer {
    max_length: f32,
 }
 impl LengthBasedDifficultyScorer {
    pub fn new() -> Self {
        Self { max_length: 100.0 }
    }
 }
 impl DifficultyScorer for LengthBasedDifficultyScorer {
    fn score<S: Sample>(&self, sample: &S) -> f32 {
        let features = sample.complexity_features();
        let length = features.get("sequence_length").copied().unwrap_or(0.0);
        (length / self.max_length).min(1.0)
    }
 }
 /// Curriculum state
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CurriculumState {
    pub step: usize,
    pub difficulty_threshold: f32,
    pub performance_history: Vec<f32>,
 }
 impl CurriculumState {
    pub fn new() -> Self {
        Self {
            step: 0,
            difficulty_threshold: 0.0,
            performance_history: Vec::new(),
        }
    }
 }
 /// Curriculum strategy trait
 pub trait CurriculumStrategy: Send + Sync {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S>;
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState);
 }
 /// Easy-to-hard strategy
 #[derive(Debug, Clone)]
 pub struct EasyToHardStrategy {
    initial_threshold: f32,
    threshold_increment: f32,
 }
 impl EasyToHardStrategy {
    pub fn new(initial_threshold: f32, threshold_increment: f32) -> Self {
        Self {
            initial_threshold,
            threshold_increment,
        }
    }
 }
 impl CurriculumStrategy for EasyToHardStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        let mut scored_samples: Vec<(S, f32)> = samples
            .iter()
            .map(|s| (s.clone(), difficulty_scorer.score(s)))
            .collect();
        // Filter by current difficulty threshold
        scored_samples.retain(|(_, score)| *score <= curriculum_state.difficulty_threshold);
        // Sort by difficulty (easiest first)
        scored_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
        // If not enough samples, include some harder ones
        if scored_samples.len() < batch_size {
            let mut all_samples: Vec<(S, f32)> = samples
                .iter()
                .map(|s| (s.clone(), difficulty_scorer.score(s)))
                .collect();
            all_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
            scored_samples = all_samples.into_iter().take(batch_size).collect();
        }
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        curriculum_state.difficulty_threshold = 
            (self.initial_threshold + curriculum_state.step as f32 * self.threshold_increment).min(1.0);
    }
 }
 /// Schedule trait
 pub trait Schedule: Send + Sync {
    fn get_difficulty_at_step(&self, step: usize) -> f32;
 }
 /// Linear schedule
 #[derive(Debug, Clone)]
 pub struct LinearSchedule {
    initial_difficulty: f32,
    final_difficulty: f32,
    total_steps: usize,
 }
 impl LinearSchedule {
    pub fn new(initial_difficulty: f32, final_difficulty: f32, total_steps: usize) -> Self {
        Self {
            initial_difficulty,
            final_difficulty,
            total_steps,
        }
    }
 }
 impl Schedule for LinearSchedule {
    fn get_difficulty_at_step(&self, step: usize) -> f32 {
        if step >= self.total_steps {
            return self.final_difficulty;
        }
        let progress = step as f32 / self.total_steps as f32;
        self.initial_difficulty + progress * (self.final_difficulty - self.initial_difficulty)
    }
 }
 fn main() -> Result<()> {
    println!("🎯 Testing RTX Curriculum Learning Implementation");
    println!("================================================");
    // Test 1: Difficulty Scoring
    println!("\n📊 Test 1: Difficulty Scoring");
    let scorer = LengthBasedDifficultyScorer::new();
    let easy_sample = TestSample::new(1, vec![1.0, 2.0], None);
    let hard_sample = TestSample::new(2, vec![1.0, 2.0, 3.0, 4.0, 5.0], None);
    let easy_score = scorer.score(&easy_sample);
    let hard_score = scorer.score(&hard_sample);
    println!("   Easy sample (len=2): difficulty = {:.3}", easy_score);
    println!("   Hard sample (len=5): difficulty = {:.3}", hard_score);
    assert!(easy_score < hard_score, "❌ Easy sample should have lower difficulty");
    assert!(easy_score >= 0.0 && easy_score <= 1.0, "❌ Easy score should be normalized");
    assert!(hard_score >= 0.0 && hard_score <= 1.0, "❌ Hard score should be normalized");
    println!("   ✅ Difficulty scoring works correctly");
    // Test 2: Curriculum Strategy
    println!("\n📈 Test 2: Easy-to-Hard Curriculum Strategy");
    let strategy = EasyToHardStrategy::new(0.1, 0.05);
    let samples = vec![
        TestSample::new(1, vec![1.0; 10], None), // Hard
        TestSample::new(2, vec![1.0; 2], None),  // Easy
        TestSample::new(3, vec![1.0; 5], None),  // Medium
    ];
    let mut state = CurriculumState::new();
    // Initially should select mostly easy samples
    let selected = strategy.select_samples(&samples, &scorer, &mut state, 2);
    println!("   Selected {} samples initially", selected.len());
    let selected_lengths: Vec<usize> = selected.iter()
        .map(|s| s.data.len())
        .collect();
    println!("   Selected sample lengths: {:?}", selected_lengths);
    assert_eq!(selected.len(), 2, "❌ Should select requested batch size");
    assert!(selected_lengths.contains(&2), "❌ Easy sample should be selected initially");
    println!("   ✅ Initial selection prefers easy samples");
    // Advance curriculum
    for i in 0..20 {
        state.step += 1;
        strategy.update_difficulty_threshold(&mut state);
        println!("   Step {}: threshold = {:.3}", i + 1, state.difficulty_threshold);
    }
    let selected_later = strategy.select_samples(&samples, &scorer, &mut state, 2);
    println!("   Selected {} samples after progression", selected_later.len());
    assert_eq!(selected_later.len(), 2, "❌ Should still select requested batch size");
    println!("   ✅ Curriculum progression works");
    // Test 3: Linear Schedule
    println!("\n⏱️  Test 3: Linear Schedule");
    let schedule = LinearSchedule::new(0.1, 0.9, 100);
    let step_0 = schedule.get_difficulty_at_step(0);
    let step_50 = schedule.get_difficulty_at_step(50);
    let step_100 = schedule.get_difficulty_at_step(100);
    let step_150 = schedule.get_difficulty_at_step(150);
    println!("   Step 0: difficulty = {:.3}", step_0);
    println!("   Step 50: difficulty = {:.3}", step_50);
    println!("   Step 100: difficulty = {:.3}", step_100);
    println!("   Step 150: difficulty = {:.3}", step_150);
    assert!((step_0 - 0.1).abs() < 1e-6, "❌ Step 0 should be initial difficulty");
    assert!((step_50 - 0.5).abs() < 1e-6, "❌ Step 50 should be halfway");
    assert!((step_100 - 0.9).abs() < 1e-6, "❌ Step 100 should be final difficulty");
    assert!((step_150 - 0.9).abs() < 1e-6, "❌ Step 150 should clamp at final difficulty");
    println!("   ✅ Linear schedule works correctly");
    // Test 4: Curriculum State Serialization
    println!("\n💾 Test 4: Curriculum State Serialization");
    let mut test_state = CurriculumState::new();
    test_state.step = 42;
    test_state.difficulty_threshold = 0.75;
    test_state.performance_history = vec![0.8, 0.85, 0.9];
    let serialized = serde_json::to_string(&test_state).unwrap();
    let deserialized: CurriculumState = serde_json::from_str(&serialized).unwrap();
    assert_eq!(deserialized.step, 42, "❌ Step should match");
    assert!((deserialized.difficulty_threshold - 0.75).abs() < 1e-6, "❌ Threshold should match");
    assert_eq!(deserialized.performance_history.len(), 3, "❌ Performance history should match");
    println!("   ✅ State serialization works correctly");
    // Test 5: Integration Test
    println!("\n🔗 Test 5: Integration Test");
    println!("   Testing complete curriculum learning pipeline...");
    let curriculum_samples = vec![
        TestSample::new(1, vec![1.0; 1], None),   // Very easy
        TestSample::new(2, vec![1.0; 3], None),   // Easy
        TestSample::new(3, vec![1.0; 5], None),   // Medium
        TestSample::new(4, vec![1.0; 8], None),   // Hard
        TestSample::new(5, vec![1.0; 12], None),  // Very hard
    ];
    let mut curriculum_state = CurriculumState::new();
    let curriculum_strategy = EasyToHardStrategy::new(0.1, 0.1);
    let curriculum_schedule = LinearSchedule::new(0.1, 0.8, 10);
    for step in 0..10 {
        curriculum_state.step = step;
        curriculum_state.difficulty_threshold = curriculum_schedule.get_difficulty_at_step(step);
        curriculum_strategy.update_difficulty_threshold(&mut curriculum_state);
        let batch = curriculum_strategy.select_samples(
            &curriculum_samples,
            &scorer,
            &mut curriculum_state,
            2
        );
        let avg_length: f32 = batch.iter()
            .map(|s| s.data.len() as f32)
            .sum::<f32>() / batch.len() as f32;
        println!("   Step {}: threshold = {:.2}, avg_batch_length = {:.1}", 
                step, curriculum_state.difficulty_threshold, avg_length);
    }
    println!("   ✅ Complete pipeline integration successful");
    println!("\n🎉 All Curriculum Learning Tests Passed!");
    println!("   ✅ Difficulty scoring functions work correctly");
    println!("   ✅ Easy-to-hard strategy implemented properly");
    println!("   ✅ Linear scheduling functions correctly");
    println!("   ✅ State serialization/deserialization works");
    println!("   ✅ Complete pipeline integration successful");
    println!("\n📋 Implementation Status:");
    println!("   🟢 Core curriculum learning framework - COMPLETE");
    println!("   🟢 Difficulty scoring strategies - COMPLETE");
    println!("   🟢 Curriculum selection strategies - COMPLETE");
    println!("   🟢 Scheduling algorithms - COMPLETE");
    println!("   🟢 Performance tracking foundation - COMPLETE");
    println!("   🟢 State management and serialization - COMPLETE");
    Ok(())
 }
 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
    use super::*;
    #[test]
    fn test_curriculum_learning_basic_functionality() {
        // This ensures our curriculum learning implementation works
        let result = main();
        assert!(result.is_ok(), "Curriculum learning tests should pass");
    }
 }
@@ -1,30 +0,0 @@
 //! Curriculum learning state.
 use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 /// State of curriculum learning process
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CurriculumState {
    pub step: usize,
    pub difficulty_threshold: f32,
    pub performance_history: Vec<f32>,
    pub metadata: HashMap<String, f32>,
 }
 impl CurriculumState {
    pub fn new() -> Self {
        Self {
            step: 0,
            difficulty_threshold: 0.0,
            performance_history: Vec::new(),
            metadata: HashMap::new(),
        }
    }
 }
 impl Default for CurriculumState {
    fn default() -> Self {
        Self::new()
    }
 }
@@ -1,414 +0,0 @@
 //! Curriculum strategy implementations.
 use std::collections::HashMap;
 use super::sample::Sample;
 use super::scorers::DifficultyScorer;
 use super::state::CurriculumState;
 /// Trait for curriculum strategies
 pub trait CurriculumStrategy: Send + Sync {
    /// Select samples based on curriculum strategy
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S>;
    /// Update difficulty threshold based on strategy
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState);
 }
 /// Easy-to-hard curriculum strategy (Bengio et al.)
 #[derive(Debug, Clone)]
 pub struct EasyToHardStrategy {
    initial_threshold: f32,
    threshold_increment: f32,
 }
 impl EasyToHardStrategy {
    pub fn new(initial_threshold: f32, threshold_increment: f32) -> Self {
        Self {
            initial_threshold,
            threshold_increment,
        }
    }
 }
 impl CurriculumStrategy for EasyToHardStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        let mut scored_samples: Vec<(S, f32)> = samples
            .iter()
            .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
            .collect();
        // Filter by current difficulty threshold
        scored_samples.retain(|(_, score)| *score <= curriculum_state.difficulty_threshold);
        // If not enough samples, include some harder ones
        if scored_samples.len() < batch_size {
            let mut all_samples: Vec<(S, f32)> = samples
                .iter()
                .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
                .collect();
            all_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
            scored_samples = all_samples.into_iter().take(batch_size).collect();
        }
        // Select batch_size samples, preferring easier ones
        scored_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        curriculum_state.difficulty_threshold =
            (self.initial_threshold + curriculum_state.step as f32 * self.threshold_increment).min(1.0);
    }
 }
 /// Anti-curriculum strategy (start with hard examples)
 #[derive(Debug, Clone)]
 pub struct AntiCurriculumStrategy {
    initial_threshold: f32,
    threshold_decrement: f32,
 }
 impl AntiCurriculumStrategy {
    pub fn new(initial_threshold: f32, threshold_decrement: f32) -> Self {
        Self {
            initial_threshold,
            threshold_decrement,
        }
    }
 }
 impl CurriculumStrategy for AntiCurriculumStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        let mut scored_samples: Vec<(S, f32)> = samples
            .iter()
            .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
            .collect();
        // Filter by current difficulty threshold (but prefer harder samples)
        scored_samples.retain(|(_, score)| *score >= curriculum_state.difficulty_threshold);
        // Sort by difficulty (hardest first)
        scored_samples.sort_by(|a, b| b.1.total_cmp(&a.1));
        // If not enough samples, include easier ones
        if scored_samples.len() < batch_size {
            let mut all_samples: Vec<(S, f32)> = samples
                .iter()
                .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
                .collect();
            all_samples.sort_by(|a, b| b.1.total_cmp(&a.1));
            scored_samples = all_samples.into_iter().take(batch_size).collect();
        }
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        curriculum_state.difficulty_threshold =
            (self.initial_threshold - curriculum_state.step as f32 * self.threshold_decrement).max(0.0);
    }
 }
 /// Self-paced learning strategy (adjusts based on performance)
 #[derive(Debug, Clone)]
 pub struct SelfPacedStrategy {
    initial_threshold: f32,
    target_performance: f32,
    performance_window_size: usize,
    performance_window: Vec<f32>,
 }
 impl SelfPacedStrategy {
    pub fn new(initial_threshold: f32, target_performance: f32, performance_window_size: usize) -> Self {
        Self {
            initial_threshold,
            target_performance,
            performance_window_size,
            performance_window: Vec::new(),
        }
    }
    pub fn update_performance(&mut self, performance: f32) {
        self.performance_window.push(performance);
        if self.performance_window.len() > self.performance_window_size {
            self.performance_window.remove(0);
        }
    }
    pub fn get_performance_window(&self) -> &[f32] {
        &self.performance_window
    }
    fn current_performance(&self) -> f32 {
        if self.performance_window.is_empty() {
            self.target_performance
        } else {
            self.performance_window.iter().sum::<f32>() / self.performance_window.len() as f32
        }
    }
 }
 impl CurriculumStrategy for SelfPacedStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        let mut scored_samples: Vec<(S, f32)> = samples
            .iter()
            .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
            .collect();
        // Adjust threshold based on current performance
        let current_perf = self.current_performance();
        let adaptive_threshold = if current_perf > self.target_performance {
            curriculum_state.difficulty_threshold * 1.1 // Increase difficulty
        } else {
            curriculum_state.difficulty_threshold * 0.9 // Decrease difficulty
        };
        // Filter by adaptive threshold
        scored_samples.retain(|(_, score)| *score <= adaptive_threshold);
        // Sort by difficulty
        scored_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
        // If not enough samples, take from all available
        if scored_samples.len() < batch_size {
            let mut all_samples: Vec<(S, f32)> = samples
                .iter()
                .map(|s| (s.clone(), difficulty_scorer.score_features(&s.complexity_features())))
                .collect();
            all_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
            scored_samples = all_samples.into_iter().take(batch_size).collect();
        }
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        let current_perf = self.current_performance();
        let adjustment = if current_perf > self.target_performance {
            0.01 // Gradually increase
        } else {
            -0.01 // Gradually decrease
        };
        curriculum_state.difficulty_threshold =
            (curriculum_state.difficulty_threshold + adjustment).clamp(0.0, 1.0);
    }
 }
 /// Competency-based curriculum strategy
 #[derive(Debug, Clone)]
 pub struct CompetencyBasedStrategy {
    competencies: HashMap<String, f32>,
 }
 impl CompetencyBasedStrategy {
    pub fn new() -> Self {
        Self {
            competencies: HashMap::new(),
        }
    }
    pub fn add_competency(&mut self, competency: &str, level: f32) {
        self.competencies.insert(competency.to_string(), level);
    }
    pub fn update_competency(&mut self, competency: &str, level: f32) {
        self.competencies.insert(competency.to_string(), level);
    }
    pub fn get_competencies(&self) -> &HashMap<String, f32> {
        &self.competencies
    }
 }
 impl Default for CompetencyBasedStrategy {
    fn default() -> Self {
        Self::new()
    }
 }
 impl CurriculumStrategy for CompetencyBasedStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        _difficulty_scorer: &dyn DifficultyScorer,
        _curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        let mut suitable_samples = Vec::new();
        for sample in samples {
            let features = sample.complexity_features();
            let mut is_suitable = true;
            // Check if sample matches current competency levels
            for (competency, &required_level) in &features {
                if competency.ends_with("_difficulty") {
                    let competency_name = competency.strip_suffix("_difficulty").unwrap_or(competency);
                    if let Some(&current_level) = self.competencies.get(competency_name) {
                        if required_level > current_level * 1.2 { // Allow 20% buffer
                            is_suitable = false;
                            break;
                        }
                    }
                }
            }
            if is_suitable {
                suitable_samples.push(sample.clone());
            }
            if suitable_samples.len() >= batch_size {
                break;
            }
        }
        // If not enough suitable samples, take any available
        while suitable_samples.len() < batch_size && suitable_samples.len() < samples.len() {
            for sample in samples {
                if !suitable_samples.iter().any(|s| s.id() == sample.id()) {
                    suitable_samples.push(sample.clone());
                    if suitable_samples.len() >= batch_size {
                        break;
                    }
                }
            }
            break;
        }
        suitable_samples
    }
    fn update_difficulty_threshold(&self, _curriculum_state: &mut CurriculumState) {
        // Competency-based doesn't use a single difficulty threshold
    }
 }
 /// Data-driven curriculum discovery strategy
 #[derive(Debug, Clone)]
 pub struct DataDrivenDiscoveryStrategy {
    cluster_count: usize,
    learning_rate: f32,
    performance_data: HashMap<usize, f32>,
    discovered_patterns: Vec<String>,
 }
 impl DataDrivenDiscoveryStrategy {
    pub fn new(cluster_count: usize, learning_rate: f32) -> Self {
        Self {
            cluster_count,
            learning_rate,
            performance_data: HashMap::new(),
            discovered_patterns: Vec::new(),
        }
    }
    pub fn add_performance_data(&mut self, sample_id: usize, performance: f32) {
        self.performance_data.insert(sample_id, performance);
        self.discover_patterns();
    }
    pub fn get_discovered_patterns(&self) -> &[String] {
        &self.discovered_patterns
    }
    fn discover_patterns(&mut self) {
        // Simple pattern discovery: find performance correlations
        if self.performance_data.len() >= self.cluster_count {
            let avg_performance: f32 = self.performance_data.values().sum::<f32>()
                / self.performance_data.len() as f32;
            let pattern = if avg_performance > 0.8 {
                "high_performance_cluster"
            } else if avg_performance > 0.5 {
                "medium_performance_cluster"
            } else {
                "low_performance_cluster"
            };
            if !self.discovered_patterns.contains(&pattern.to_string()) {
                self.discovered_patterns.push(pattern.to_string());
            }
        }
    }
 }
 impl CurriculumStrategy for DataDrivenDiscoveryStrategy {
    fn select_samples<S: Sample>(
        &self,
        samples: &[S],
        difficulty_scorer: &dyn DifficultyScorer,
        _curriculum_state: &mut CurriculumState,
        batch_size: usize,
    ) -> Vec<S> {
        // Use discovered patterns to guide selection
        let mut scored_samples: Vec<(S, f32)> = samples
            .iter()
            .map(|s| {
                let base_score = difficulty_scorer.score_features(&s.complexity_features());
                let performance_adjustment = self.performance_data
                    .get(&s.id())
                    .map(|&perf| 1.0 - perf)
                    .unwrap_or(0.0);
                (s.clone(), base_score + performance_adjustment * 0.1)
            })
            .collect();
        // Sort by adjusted score
        scored_samples.sort_by(|a, b| a.1.total_cmp(&b.1));
        scored_samples
            .into_iter()
            .take(batch_size)
            .map(|(sample, _)| sample)
            .collect()
    }
    fn update_difficulty_threshold(&self, curriculum_state: &mut CurriculumState) {
        // Use discovered patterns to adjust threshold
        let avg_performance = if !self.performance_data.is_empty() {
            self.performance_data.values().sum::<f32>() / self.performance_data.len() as f32
        } else {
            0.5
        };
        curriculum_state.difficulty_threshold = avg_performance * self.learning_rate +
            curriculum_state.difficulty_threshold * (1.0 - self.learning_rate);
    }
 }
@@ -1,65 +0,0 @@
 //! Performance tracking for curriculum learning.
 use std::collections::HashMap;
 /// Performance tracking for curriculum learning
 #[derive(Debug, Clone)]
 pub struct PerformanceTracker {
    sample_performance: HashMap<usize, f32>,
    overall_performance_window: Vec<f32>,
    window_size: usize,
    #[allow(dead_code)]
    smoothing_factor: f32,
 }
 impl PerformanceTracker {
    pub fn new(window_size: usize, smoothing_factor: f32) -> Self {
        Self {
            sample_performance: HashMap::new(),
            overall_performance_window: Vec::new(),
            window_size,
            smoothing_factor,
        }
    }
    pub fn record_sample_performance(&mut self, sample_id: usize, performance: f32) {
        self.sample_performance.insert(sample_id, performance);
    }
    pub fn record_overall_performance(&mut self, performance: f32) {
        self.overall_performance_window.push(performance);
        if self.overall_performance_window.len() > self.window_size {
            self.overall_performance_window.remove(0);
        }
    }
    pub fn get_sample_performance(&self, sample_id: usize) -> Option<f32> {
        self.sample_performance.get(&sample_id).copied()
    }
    pub fn get_overall_performance(&self) -> f32 {
        if self.overall_performance_window.is_empty() {
            0.0
        } else {
            self.overall_performance_window.iter().sum::<f32>() / self.overall_performance_window.len() as f32
        }
    }
    pub fn get_performance_window(&self) -> &[f32] {
        &self.overall_performance_window
    }
    pub fn get_performance_trend(&self) -> f32 {
        if self.overall_performance_window.len() < 2 {
            return 0.0;
        }
        let recent_half = &self.overall_performance_window[self.overall_performance_window.len()/2..];
        let early_half = &self.overall_performance_window[..self.overall_performance_window.len()/2];
        let recent_avg = recent_half.iter().sum::<f32>() / recent_half.len() as f32;
        let early_avg = early_half.iter().sum::<f32>() / early_half.len() as f32;
        recent_avg - early_avg
    }
 }
@@ -1,175 +0,0 @@
 //! Error types for the RTX Transformers crate
 use thiserror::Error;
 /// Result type alias for transformer operations
 pub type Result<T> = std::result::Result<T, TransformerError>;
 /// Comprehensive error types for transformer operations
 #[derive(Error, Debug)]
 pub enum TransformerError {
    /// Tensor operation errors
    #[error("Tensor operation failed: {0}")]
    TensorError(#[from] rtx_tensor::TensorError),
    /// Autograd computation errors
    #[error("Gradient computation failed: {0}")]
    AutogradError(#[from] rtx_autograd::AutogradError),
    /// Runtime/CUDA errors
    #[error("Runtime error: {0}")]
    RuntimeError(#[from] rtx_runtime::RuntimeError),
    /// Optimizer configuration errors
    #[error("Optimizer configuration error: {message}")]
    OptimizerError { message: String },
    /// Learning rate scheduler errors
    #[error("Learning rate scheduler error: {message}")]
    SchedulerError { message: String },
    /// Layer configuration errors
    #[error("Layer configuration error: {message}")]
    LayerError { message: String },
    /// Architecture configuration errors
    #[error("Architecture configuration error: {message}")]
    ArchitectureError { message: String },
    /// Training configuration errors
    #[error("Training configuration error: {message}")]
    TrainingError { message: String },
    /// Tokenization errors
    #[error("Tokenization error: {message}")]
    TokenizationError { message: String },
    /// Revolutionary integration errors
    #[error("Revolutionary integration error: {message}")]
    RevolutionaryError { message: String },
    /// Invalid parameter errors
    #[error("Invalid parameter: {parameter} = {value}, reason: {reason}")]
    InvalidParameter {
        parameter: String,
        value: String,
        reason: String,
    },
    /// Shape mismatch errors
    #[error("Shape mismatch: expected {expected:?}, got {actual:?}")]
    ShapeMismatch {
        expected: Vec<usize>,
        actual: Vec<usize>,
    },
    /// Dimension errors
    #[error("Dimension error: {message}")]
    DimensionError { message: String },
    /// Memory allocation errors
    #[error("Memory allocation error: {message}")]
    MemoryError { message: String },
    /// Convergence errors
    #[error("Convergence error: {message}")]
    ConvergenceError { message: String },
    /// I/O errors
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),
    /// Serialization errors
    #[error("Serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),
    /// Generic errors
    #[error("Generic error: {0}")]
    Generic(#[from] anyhow::Error),
 }
 impl TransformerError {
    /// Create a new optimizer error
    pub fn optimizer<S: Into<String>>(message: S) -> Self {
        Self::OptimizerError {
            message: message.into(),
        }
    }
    /// Create a new scheduler error
    pub fn scheduler<S: Into<String>>(message: S) -> Self {
        Self::SchedulerError {
            message: message.into(),
        }
    }
    /// Create a new layer error
    pub fn layer<S: Into<String>>(message: S) -> Self {
        Self::LayerError {
            message: message.into(),
        }
    }
    /// Create a new architecture error
    pub fn architecture<S: Into<String>>(message: S) -> Self {
        Self::ArchitectureError {
            message: message.into(),
        }
    }
    /// Create a new training error
    pub fn training<S: Into<String>>(message: S) -> Self {
        Self::TrainingError {
            message: message.into(),
        }
    }
    /// Create a new tokenization error
    pub fn tokenization<S: Into<String>>(message: S) -> Self {
        Self::TokenizationError {
            message: message.into(),
        }
    }
    /// Create a new revolutionary integration error
    pub fn revolutionary<S: Into<String>>(message: S) -> Self {
        Self::RevolutionaryError {
            message: message.into(),
        }
    }
    /// Create a new invalid parameter error
    pub fn invalid_parameter<S: Into<String>>(parameter: S, value: S, reason: S) -> Self {
        Self::InvalidParameter {
            parameter: parameter.into(),
            value: value.into(),
            reason: reason.into(),
        }
    }
    /// Create a new shape mismatch error
    pub fn shape_mismatch(expected: Vec<usize>, actual: Vec<usize>) -> Self {
        Self::ShapeMismatch { expected, actual }
    }
    /// Create a new dimension error
    pub fn dimension<S: Into<String>>(message: S) -> Self {
        Self::DimensionError {
            message: message.into(),
        }
    }
    /// Create a new memory error
    pub fn memory<S: Into<String>>(message: S) -> Self {
        Self::MemoryError {
            message: message.into(),
        }
    }
    /// Create a new convergence error
    pub fn convergence<S: Into<String>>(message: S) -> Self {
        Self::ConvergenceError {
            message: message.into(),
        }
    }
 }
@@ -1,29 +0,0 @@
 //! Error types for the RTX Transformers crate (Minimal version)
 use thiserror::Error;
 /// Result type alias for transformer operations
 pub type Result<T> = std::result::Result<T, TransformerError>;
 /// Minimal error types for transformer operations
 #[derive(Error, Debug)]
 pub enum TransformerError {
    /// Generic errors
    #[error("Transformer error: {0}")]
    Generic(String),
    /// I/O errors
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),
    /// Anyhow errors
    #[error("Error: {0}")]
    Anyhow(#[from] anyhow::Error),
 }
 impl TransformerError {
    /// Create a generic error
    pub fn generic<S: Into<String>>(message: S) -> Self {
        Self::Generic(message.into())
    }
 }
@@ -1,447 +0,0 @@
 //! Graph Attention implementation for attention mechanisms on graph-structured data
 //!
 //! Provides multiple graph attention mechanisms:
 //! - Standard Graph Attention (GAT)
 //! - Edge-aware Graph Attention  
 //! - Multi-head Graph Attention
 //! - Gated Graph Attention
 use crate::{Result, TransformerError};
 use crate::graph::{GraphBatch, Graph, GraphAttentionType};
 use rtx_tensor::{Tensor, Device};
 use serde::{Deserialize, Serialize};
 /// Graph attention mechanism configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GraphAttentionConfig {
    /// Input feature dimension
    pub input_dim: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Type of attention mechanism
    pub attention_type: GraphAttentionType,
    /// Edge feature dimension (for edge-aware attention)
    pub edge_dim: Option<usize>,
    /// Whether to concatenate or average multi-head outputs
    pub concat_heads: bool,
    /// Dropout probability
    pub dropout: f32,
    /// Activation function for attention
    pub attention_activation: String,
    /// Gate activation function (for gated attention)
    pub gate_activation: String,
    /// Whether to use bias in linear layers
    pub use_bias: bool,
    /// Attention temperature scaling
    pub temperature: f32,
 }
 impl GraphAttentionConfig {
    /// Create new graph attention configuration
    pub fn new(input_dim: usize, num_heads: usize, attention_type: GraphAttentionType) -> Self {
        Self {
            input_dim,
            num_heads,
            attention_type,
            edge_dim: None,
            concat_heads: true,
            dropout: 0.1,
            attention_activation: "leaky_relu".to_string(),
            gate_activation: "sigmoid".to_string(),
            use_bias: true,
            temperature: 1.0,
        }
    }
    /// Set edge dimension for edge-aware attention
    pub fn with_edge_dim(mut self, edge_dim: usize) -> Self {
        self.edge_dim = Some(edge_dim);
        self
    }
    /// Set whether to concatenate heads
    pub fn with_concat_heads(mut self, concat_heads: bool) -> Self {
        self.concat_heads = concat_heads;
        self
    }
    /// Set gate activation function
    pub fn with_gate_activation(mut self, activation: &str) -> Self {
        self.gate_activation = activation.to_string();
        self
    }
 }
 /// Output from graph attention layer
 #[derive(Debug)]
 pub struct AttentionOutput {
    /// Updated node features
    pub node_features: Tensor,
    /// Updated edge features (if edge-aware)
    pub edge_features: Option<Tensor>,
    /// Attention weights [num_heads, num_nodes, num_nodes]
    pub attention_weights: Option<Tensor>,
    /// Gate values (if gated attention)
    pub gate_values: Option<Tensor>,
 }
 /// Graph attention mechanism types
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum AttentionMechanism {
    /// Dot-product attention
    DotProduct,
    /// Additive attention
    Additive,
    /// Scaled dot-product attention
    ScaledDotProduct,
 }
 /// Graph attention layer implementation
 pub struct GraphAttention {
    config: GraphAttentionConfig,
    device: Device,
    // Attention parameters
    query_projection: Tensor,    // [input_dim, head_dim * num_heads]
    key_projection: Tensor,      // [input_dim, head_dim * num_heads]
    value_projection: Tensor,    // [input_dim, head_dim * num_heads]
    // Edge attention parameters (for edge-aware)
    edge_query_projection: Option<Tensor>,
    edge_key_projection: Option<Tensor>,
    edge_value_projection: Option<Tensor>,
    // Output projection
    output_projection: Tensor,   // [head_dim * num_heads, input_dim]
    // Gate parameters (for gated attention)
    gate_projection: Option<Tensor>,
    // Attention mechanism
    attention_mechanism: AttentionMechanism,
    // Head dimension
    head_dim: usize,
 }
 impl GraphAttention {
    /// Create a new graph attention layer
    pub fn new(config: GraphAttentionConfig, device: &Device) -> Result<Self> {
        if config.input_dim % config.num_heads != 0 {
            return Err(TransformerError::InvalidInput(format!(
                "Input dimension {} must be divisible by number of heads {}",
                config.input_dim, config.num_heads
            )));
        }
        let head_dim = config.input_dim / config.num_heads;
        let total_dim = head_dim * config.num_heads;
        // Initialize query, key, value projections
        let query_projection = Tensor::randn(&[config.input_dim, total_dim], device)?;
        let key_projection = Tensor::randn(&[config.input_dim, total_dim], device)?;
        let value_projection = Tensor::randn(&[config.input_dim, total_dim], device)?;
        // Initialize edge projections if edge-aware
        let (edge_query_projection, edge_key_projection, edge_value_projection) = 
            if config.attention_type == GraphAttentionType::EdgeAware {
                if let Some(edge_dim) = config.edge_dim {
                    (
                        Some(Tensor::randn(&[edge_dim, total_dim], device)?),
                        Some(Tensor::randn(&[edge_dim, total_dim], device)?),
                        Some(Tensor::randn(&[edge_dim, total_dim], device)?),
                    )
                } else {
                    return Err(TransformerError::InvalidInput(
                        "Edge dimension must be specified for edge-aware attention".to_string()
                    ));
                }
            } else {
                (None, None, None)
            };
        // Output projection
        let output_projection = if config.concat_heads {
            Tensor::randn(&[total_dim, config.input_dim], device)?
        } else {
            Tensor::randn(&[head_dim, config.input_dim], device)?
        };
        // Gate projection for gated attention
        let gate_projection = if config.attention_type == GraphAttentionType::Gated {
            Some(Tensor::randn(&[config.input_dim, config.input_dim], device)?)
        } else {
            None
        };
        let attention_mechanism = match config.attention_type {
            GraphAttentionType::Standard | GraphAttentionType::EdgeAware => AttentionMechanism::DotProduct,
            GraphAttentionType::MultiHead => AttentionMechanism::ScaledDotProduct,
            GraphAttentionType::Gated => AttentionMechanism::Additive,
        };
        Ok(Self {
            config,
            device: device.clone(),
            query_projection,
            key_projection,
            value_projection,
            edge_query_projection,
            edge_key_projection,
            edge_value_projection,
            output_projection,
            gate_projection,
            attention_mechanism,
            head_dim,
        })
    }
    /// Get the configuration
    pub fn config(&self) -> &GraphAttentionConfig {
        &self.config
    }
    /// Compute attention scores between nodes
    fn compute_attention_scores(
        &self, 
        queries: &Tensor, 
        keys: &Tensor, 
        edge_indices: &Tensor,
        edge_features: Option<&Tensor>,
    ) -> Result<Tensor> {
        let num_edges = edge_indices.shape().dims()[1];
        let num_heads = self.config.num_heads;
        // Get source and target node indices
        let source_indices = edge_indices.get(&[0])?;  // [num_edges]
        let target_indices = edge_indices.get(&[1])?;  // [num_edges]
        // Gather queries and keys for edges
        let source_queries = queries.index_select(0, &source_indices)?;  // [num_edges, head_dim * num_heads]
        let target_keys = keys.index_select(0, &target_indices)?;        // [num_edges, head_dim * num_heads]
        // Reshape for multi-head attention: [num_edges, num_heads, head_dim]
        let source_queries = source_queries.view([num_edges, num_heads, self.head_dim])?;
        let target_keys = target_keys.view([num_edges, num_heads, self.head_dim])?;
        // Compute attention scores
        let mut attention_scores = match self.attention_mechanism {
            AttentionMechanism::DotProduct => {
                // Simple dot product: sum over head dimension
                source_queries.mul(&target_keys)?.sum_dim(&[2], false)?
            }
            AttentionMechanism::ScaledDotProduct => {
                // Scaled dot product
                let scores = source_queries.mul(&target_keys)?.sum_dim(&[2], false)?;
                let scale = (self.head_dim as f32).sqrt();
                scores.div(&Tensor::full(&[], (scale) as f32, &&self.device)?)?
            }
            AttentionMechanism::Additive => {
                // Additive attention (simplified)
                let combined = source_queries.add(&target_keys)?;
                combined.sum_dim(&[2], false)?
            }
        };
        // Add edge features if edge-aware attention
        if self.config.attention_type == GraphAttentionType::EdgeAware {
            if let (Some(edge_feats), Some(edge_q_proj)) = (edge_features, &self.edge_query_projection) {
                let edge_contributions = edge_feats.matmul(edge_q_proj)?;  // [num_edges, head_dim * num_heads]
                let edge_contributions = edge_contributions.view([num_edges, num_heads, self.head_dim])?;
                let edge_scores = edge_contributions.sum_dim(&[2], false)?;  // [num_edges, num_heads]
                attention_scores = attention_scores.add(&edge_scores)?;
            }
        }
        // Apply temperature scaling
        if self.config.temperature != 1.0 {
            attention_scores = attention_scores.div(&Tensor::full(&[], (self.config.temperature) as f32, &&self.device)?)?;
        }
        // Apply activation (leaky_relu is default)
        attention_scores = self.apply_activation(&attention_scores, &self.config.attention_activation)?;
        Ok(attention_scores)
    }
    /// Apply softmax attention over edges for each node
    fn apply_attention_softmax(&self, attention_scores: &Tensor, edge_indices: &Tensor) -> Result<Tensor> {
        let num_edges = edge_indices.shape().dims()[1];
        let num_heads = self.config.num_heads;
        let source_indices = edge_indices.get(&[0])?;
        // For each source node, apply softmax over its outgoing edges
        // This is a simplified implementation - full implementation would use scatter_softmax
        // Apply softmax directly (simplified - assumes proper grouping)
        let attention_weights = attention_scores.softmax(-1)?;  // [num_edges, num_heads]
        Ok(attention_weights)
    }
    /// Apply activation function
    fn apply_activation(&self, input: &Tensor, activation: &str) -> Result<Tensor> {
        match activation {
            "relu" => input.relu(),
            "leaky_relu" => {
                // Simplified leaky_relu with alpha=0.2
                let alpha = Tensor::full(&[], 0.2, &self.device);
                let positive = input.relu()?;
                let negative = input.clamp(None, Some(0.0))?.mul(&alpha)?;
                positive.add(&negative)
            }
            "sigmoid" => input.sigmoid(),
            "tanh" => input.tanh(),
            _ => input.clone(),
        }
    }
 }
 impl GraphLayer for GraphAttention {
    fn forward(&self, graph: &GraphBatch) -> Result<AttentionOutput> {
        let num_nodes = graph.num_nodes();
        let num_edges = graph.num_edges();
        if num_nodes == 0 {
            return Err(TransformerError::InvalidInput(
                "Cannot process empty graph".to_string()
            ));
        }
        // Project node features to queries, keys, values
        let queries = graph.node_features.matmul(&self.query_projection)?;  // [num_nodes, head_dim * num_heads]
        let keys = graph.node_features.matmul(&self.key_projection)?;       // [num_nodes, head_dim * num_heads]
        let values = graph.node_features.matmul(&self.value_projection)?;   // [num_nodes, head_dim * num_heads]
        // Compute attention scores
        let attention_scores = self.compute_attention_scores(
            &queries, 
            &keys, 
            &graph.edge_indices, 
            Some(&graph.edge_features)
        )?;
        // Apply softmax to get attention weights
        let attention_weights = self.apply_attention_softmax(&attention_scores, &graph.edge_indices)?;
        // Apply attention to values
        let source_indices = graph.edge_indices.get(&[0])?;
        let target_indices = graph.edge_indices.get(&[1])?;
        // Gather values for target nodes
        let target_values = values.index_select(0, &target_indices)?;  // [num_edges, head_dim * num_heads]
        let target_values = target_values.view([num_edges, self.config.num_heads, self.head_dim])?;
        // Weight values by attention
        let attention_weights_expanded = attention_weights.unsqueeze(2)?;  // [num_edges, num_heads, 1]
        let weighted_values = target_values.mul(&attention_weights_expanded)?;  // [num_edges, num_heads, head_dim]
        // Aggregate weighted values for each source node
        // Simplified aggregation - in full implementation would use scatter_add
        let aggregated_values = if self.config.concat_heads {
            weighted_values.view([num_edges, self.config.num_heads * self.head_dim])?
        } else {
            weighted_values.mean_dim(&[1], false)?  // Average over heads
        };
        // Create output tensor initialized with input features
        let mut output_features = graph.node_features.clone();
        // In a full implementation, we would scatter_add the aggregated values
        // For now, we'll apply a simple linear transformation as a placeholder
        output_features = output_features.matmul(&self.output_projection)?;
        // Apply gating if enabled
        let gate_values = if self.config.attention_type == GraphAttentionType::Gated {
            if let Some(ref gate_proj) = self.gate_projection {
                let gates = graph.node_features.matmul(gate_proj)?;
                let gate_vals = self.apply_activation(&gates, &self.config.gate_activation)?;
                output_features = output_features.mul(&gate_vals)?;
                Some(gate_vals)
            } else {
                None
            }
        } else {
            None
        };
        // Process edge features if edge-aware
        let updated_edge_features = if self.config.attention_type == GraphAttentionType::EdgeAware {
            if let Some(ref edge_v_proj) = self.edge_value_projection {
                Some(graph.edge_features.matmul(edge_v_proj)?)
            } else {
                None
            }
        } else {
            None
        };
        Ok(AttentionOutput {
            node_features: output_features,
            edge_features: updated_edge_features,
            attention_weights: Some(attention_weights),
            gate_values,
        })
    }
    fn layer_type(&self) -> &'static str {
        match self.config.attention_type {
            GraphAttentionType::Standard => "StandardGraphAttention",
            GraphAttentionType::EdgeAware => "EdgeAwareGraphAttention", 
            GraphAttentionType::MultiHead => "MultiHeadGraphAttention",
            GraphAttentionType::Gated => "GatedGraphAttention",
        }
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        let mut params = vec![
            &self.query_projection,
            &self.key_projection, 
            &self.value_projection,
            &self.output_projection,
        ];
        if let Some(ref edge_q) = self.edge_query_projection {
            params.push(edge_q);
        }
        if let Some(ref edge_k) = self.edge_key_projection {
            params.push(edge_k);
        }
        if let Some(ref edge_v) = self.edge_value_projection {
            params.push(edge_v);
        }
        if let Some(ref gate) = self.gate_projection {
            params.push(gate);
        }
        params
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![
            &mut self.query_projection,
            &mut self.key_projection,
            &mut self.value_projection, 
            &mut self.output_projection,
        ];
        if let Some(ref mut edge_q) = self.edge_query_projection {
            params.push(edge_q);
        }
        if let Some(ref mut edge_k) = self.edge_key_projection {
            params.push(edge_k);
        }
        if let Some(ref mut edge_v) = self.edge_value_projection {
            params.push(edge_v);
        }
        if let Some(ref mut gate) = self.gate_projection {
            params.push(gate);
        }
        params
    }
 }
@@ -1,595 +0,0 @@
 //! Graph Pooling implementations for creating graph-level representations
 //!
 //! Provides multiple strategies for pooling node-level features into 
 //! graph-level representations:
 //! - Global pooling (mean, max, sum)
 //! - Hierarchical graph pooling
 //! - Set2Set pooling
 //! - Attention-based pooling
 use crate::{Result, TransformerError};
 use crate::graph::{GraphBatch, PoolingStrategy};
 use rtx_tensor::{Tensor, Device};
 /// Graph pooling trait for converting node features to graph representations
 pub trait GraphPooling: Send + Sync {
    /// Forward pass to pool node features into graph-level representation
    fn forward(&self, graph: &GraphBatch) -> Result<Tensor>;
    /// Get the strategy used by this pooling layer
    fn strategy(&self) -> &PoolingStrategy;
    /// Get input dimension
    fn input_dim(&self) -> usize;
    /// Get device
    fn device(&self) -> &Device;
    /// Get parameters
    fn parameters(&self) -> Vec<&Tensor>;
    /// Get mutable parameters
    fn parameters_mut(&mut self) -> Vec<&mut Tensor>;
 }
 /// Global pooling methods
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum GlobalPoolingType {
    Mean,
    Max, 
    Sum,
    Attention,
 }
 /// Global graph pooling implementation
 pub struct GlobalPooling {
    strategy: PoolingStrategy,
    input_dim: usize,
    device: Device,
    pooling_type: GlobalPoolingType,
    // Attention pooling parameters (if using attention)
    attention_weights: Option<Tensor>,
 }
 impl GlobalPooling {
    /// Create new global pooling layer
    pub fn new(input_dim: usize, device: &Device) -> Result<Self> {
        Ok(Self {
            strategy: PoolingStrategy::Global,
            input_dim,
            device: device.clone(),
            pooling_type: GlobalPoolingType::Mean,
            attention_weights: None,
        })
    }
    /// Create global pooling with specific type
    pub fn with_type(input_dim: usize, pooling_type: GlobalPoolingType, device: &Device) -> Result<Self> {
        let attention_weights = if pooling_type == GlobalPoolingType::Attention {
            Some(Tensor::randn(&[input_dim, 1], device)?)
        } else {
            None
        };
        Ok(Self {
            strategy: PoolingStrategy::Global,
            input_dim,
            device: device.clone(),
            pooling_type,
            attention_weights,
        })
    }
 }
 impl GraphPooling for GlobalPooling {
    fn forward(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_graphs = graph.batch_size;
        let node_features = &graph.node_features;
        // Pool features for each graph in the batch
        let mut graph_representations = Vec::new();
        for i in 0..num_graphs {
            let start_idx = graph.graph_boundaries[i];
            let end_idx = graph.graph_boundaries[i + 1];
            if start_idx >= end_idx {
                // Handle empty graph case
                let empty_repr = Tensor::zeros([1, self.input_dim], &self.device)?;
                graph_representations.push(empty_repr);
                continue;
            }
            // Extract nodes for this graph
            let graph_nodes = node_features.narrow(0, start_idx, end_idx - start_idx)?;
            // Apply pooling
            let pooled = match self.pooling_type {
                GlobalPoolingType::Mean => {
                    graph_nodes.mean_dim(&[0], true)?
                }
                GlobalPoolingType::Max => {
                    graph_nodes.max_dim(0, true)?.0
                }
                GlobalPoolingType::Sum => {
                    graph_nodes.sum_dim(&[0], true)?
                }
                GlobalPoolingType::Attention => {
                    if let Some(ref attention_weights) = self.attention_weights {
                        // Compute attention scores
                        let scores = graph_nodes.matmul(attention_weights)?;  // [num_nodes, 1]
                        let attention = scores.softmax(0)?;  // [num_nodes, 1]
                        // Weighted sum
                        let weighted = graph_nodes.mul(&attention)?;
                        weighted.sum_dim(&[0], true)?
                    } else {
                        return Err(TransformerError::InvalidInput(
                            "Attention weights not initialized".to_string()
                        ));
                    }
                }
            };
            graph_representations.push(pooled);
        }
        // Stack all graph representations
        if graph_representations.is_empty() {
            Tensor::zeros([0, self.input_dim], &self.device)
        } else {
            Tensor::cat(&graph_representations, 0)  // [batch_size, input_dim]
        }
    }
    fn strategy(&self) -> &PoolingStrategy {
        &self.strategy
    }
    fn input_dim(&self) -> usize {
        self.input_dim
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        if let Some(ref weights) = self.attention_weights {
            vec![weights]
        } else {
            vec![]
        }
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        if let Some(ref mut weights) = self.attention_weights {
            vec![weights]
        } else {
            vec![]
        }
    }
 }
 /// Hierarchical graph pooling implementation
 pub struct HierarchicalPooling {
    strategy: PoolingStrategy,
    input_dim: usize,
    device: Device,
    num_levels: usize,
    pooling_ratio: f32,
    // Learnable parameters for node selection
    node_projection: Tensor,        // [input_dim, 1] for computing node scores
    edge_projection: Option<Tensor>, // [edge_dim, 1] for edge scores
 }
 impl HierarchicalPooling {
    /// Create new hierarchical pooling layer
    pub fn new(
        input_dim: usize, 
        num_levels: usize, 
        pooling_ratio: f32,
        device: &Device
    ) -> Result<Self> {
        let node_projection = Tensor::randn(&[input_dim, 1], device)?;
        Ok(Self {
            strategy: PoolingStrategy::Hierarchical,
            input_dim,
            device: device.clone(),
            num_levels,
            pooling_ratio,
            node_projection,
            edge_projection: None,
        })
    }
    /// Select top-k nodes based on learned scores
    fn select_nodes(&self, node_features: &Tensor, k: usize) -> Result<(Tensor, Tensor)> {
        // Compute node scores
        let scores = node_features.matmul(&self.node_projection)?;  // [num_nodes, 1]
        let scores = scores.squeeze(Some(1))?;  // [num_nodes]
        // Get top-k nodes
        let (top_values, top_indices) = scores.topk(k, 0, true, true)?;
        // Select corresponding node features
        let selected_features = node_features.index_select(0, &top_indices)?;
        Ok((selected_features, top_indices))
    }
    /// Coarsen graph by clustering nodes
    fn coarsen_graph(&self, graph: &GraphBatch, selected_indices: &Tensor) -> Result<GraphBatch> {
        // Simplified coarsening - in practice would properly handle edge connections
        let selected_nodes = graph.node_features.index_select(0, selected_indices)?;
        let num_selected = selected_indices.shape().dims()[0];
        // Create simplified edge structure (placeholder)
        let new_edge_indices = if num_selected > 1 {
            // Connect all selected nodes in a chain (simplified)
            let mut edge_list = Vec::new();
            for i in 0..(num_selected - 1) {
                edge_list.push(i as i64);
                edge_list.push((i + 1) as i64);
            }
            Tensor::from_vec(edge_list, [2, num_selected - 1], &self.device)?
        } else {
            Tensor::zeros([2, 0], &self.device)?
        };
        let new_edge_features = Tensor::zeros([new_edge_indices.shape().dims()[1], graph.edge_features.shape().dims()[1]], &self.device)?;
        GraphBatch::new(
            selected_nodes,
            new_edge_features,
            new_edge_indices,
            vec![0, num_selected],  // Single graph
            self.device.clone(),
        )
    }
 }
 impl GraphPooling for HierarchicalPooling {
    fn forward(&self, graph: &GraphBatch) -> Result<Tensor> {
        let mut current_graph = graph.clone();
        let mut representations = Vec::new();
        // Apply hierarchical pooling for each level
        for level in 0..self.num_levels {
            let num_graphs = current_graph.batch_size;
            let mut level_representations = Vec::new();
            for i in 0..num_graphs {
                let start_idx = current_graph.graph_boundaries[i];
                let end_idx = current_graph.graph_boundaries[i + 1];
                if start_idx >= end_idx {
                    let empty_repr = Tensor::zeros([1, self.input_dim], &self.device)?;
                    level_representations.push(empty_repr);
                    continue;
                }
                let graph_nodes = current_graph.node_features.narrow(0, start_idx, end_idx - start_idx)?;
                let num_nodes = end_idx - start_idx;
                // Determine number of nodes to select
                let k = ((num_nodes as f32 * self.pooling_ratio).ceil() as usize).max(1);
                // Select top-k nodes
                let (selected_features, _indices) = self.select_nodes(&graph_nodes, k)?;
                // Global pooling of selected nodes
                let pooled = selected_features.mean_dim(&[0], true)?;
                level_representations.push(pooled);
            }
            // Store representations from this level
            if !level_representations.is_empty() {
                let level_repr = Tensor::cat(&level_representations, 0)?;
                representations.push(level_repr);
            }
            // Create coarsened graph for next level (if not last level)
            if level < self.num_levels - 1 {
                // Simplified: just use current graph (in practice would coarsen)
                // current_graph = coarsen_graph(current_graph)?;
            }
        }
        // Combine representations from all levels
        if representations.is_empty() {
            Tensor::zeros([graph.batch_size, self.input_dim], &self.device)
        } else {
            // For simplicity, return the last level's representation
            representations.into_iter().last().unwrap()
        }
    }
    fn strategy(&self) -> &PoolingStrategy {
        &self.strategy
    }
    fn input_dim(&self) -> usize {
        self.input_dim
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        let mut params = vec![&self.node_projection];
        if let Some(ref edge_proj) = self.edge_projection {
            params.push(edge_proj);
        }
        params
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![&mut self.node_projection];
        if let Some(ref mut edge_proj) = self.edge_projection {
            params.push(edge_proj);
        }
        params
    }
 }
 /// Set2Set pooling implementation
 pub struct Set2SetPooling {
    strategy: PoolingStrategy,
    input_dim: usize,
    device: Device,
    num_iterations: usize,
    // LSTM-like parameters for Set2Set
    lstm_input_size: usize,
    lstm_hidden_size: usize,
    lstm_weights: Tensor,
    lstm_bias: Tensor,
    attention_projection: Tensor,
 }
 impl Set2SetPooling {
    /// Create new Set2Set pooling layer
    pub fn new(input_dim: usize, device: &Device) -> Result<Self> {
        let num_iterations = 3;
        let lstm_hidden_size = input_dim;
        let lstm_input_size = 2 * input_dim;  // Memory vector + pooled representation
        let lstm_weights = Tensor::randn(&[lstm_input_size + lstm_hidden_size, 4 * lstm_hidden_size], device)?;
        let lstm_bias = Tensor::zeros([4 * lstm_hidden_size], device)?;
        let attention_projection = Tensor::randn(&[lstm_hidden_size, 1], device)?;
        Ok(Self {
            strategy: PoolingStrategy::Set2Set,
            input_dim,
            device: device.clone(),
            num_iterations,
            lstm_input_size,
            lstm_hidden_size,
            lstm_weights,
            lstm_bias,
            attention_projection,
        })
    }
    /// Apply LSTM cell
    fn lstm_cell(&self, input: &Tensor, hidden: &Tensor, cell: &Tensor) -> Result<(Tensor, Tensor)> {
        // Concatenate input and hidden state
        let combined = Tensor::cat(&[input, hidden], -1)?;
        // Apply linear transformation
        let gates = combined.matmul(&self.lstm_weights)?.add(&self.lstm_bias)?;
        // Split into 4 gates
        let gate_size = self.lstm_hidden_size;
        let input_gate = gates.narrow(-1, 0, gate_size)?.sigmoid()?;
        let forget_gate = gates.narrow(-1, gate_size, gate_size)?.sigmoid()?;
        let cell_gate = gates.narrow(-1, 2 * gate_size, gate_size)?.tanh()?;
        let output_gate = gates.narrow(-1, 3 * gate_size, gate_size)?.sigmoid()?;
        // Update cell state
        let new_cell = forget_gate.mul(cell)?.add(&input_gate.mul(&cell_gate)?)?;
        // Update hidden state
        let new_hidden = output_gate.mul(&new_cell.tanh()?)?;
        Ok((new_hidden, new_cell))
    }
    /// Compute attention weights
    fn compute_attention(&self, query: &Tensor, node_features: &Tensor) -> Result<Tensor> {
        // Project query to attention space
        let attention_scores = query.matmul(&self.attention_projection)?;  // [1, 1]
        // Compute attention over all nodes (simplified)
        let num_nodes = node_features.shape().dims()[0];
        let scores = attention_scores.expand(&[num_nodes, 1])?;
        let attention_weights = scores.softmax(0)?;
        Ok(attention_weights)
    }
 }
 impl GraphPooling for Set2SetPooling {
    fn forward(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_graphs = graph.batch_size;
        let mut graph_representations = Vec::new();
        for i in 0..num_graphs {
            let start_idx = graph.graph_boundaries[i];
            let end_idx = graph.graph_boundaries[i + 1];
            if start_idx >= end_idx {
                let empty_repr = Tensor::zeros([1, 2 * self.input_dim], &self.device)?;
                graph_representations.push(empty_repr);
                continue;
            }
            let graph_nodes = graph.node_features.narrow(0, start_idx, end_idx - start_idx)?;
            let num_nodes = end_idx - start_idx;
            // Initialize LSTM states
            let mut hidden = Tensor::zeros([1, self.lstm_hidden_size], &self.device)?;
            let mut cell = Tensor::zeros([1, self.lstm_hidden_size], &self.device)?;
            let mut memory = Tensor::zeros([1, self.input_dim], &self.device)?;
            // Run Set2Set iterations
            for _iter in 0..self.num_iterations {
                // Compute attention
                let attention_weights = self.compute_attention(&hidden, &graph_nodes)?;  // [num_nodes, 1]
                // Compute attended representation
                let attended = graph_nodes.mul(&attention_weights)?.sum_dim(&[0], true)?;  // [1, input_dim]
                // Update memory
                memory = memory.add(&attended)?;
                // Prepare LSTM input
                let lstm_input = Tensor::cat(&[memory.clone(), attended], -1)?;  // [1, 2 * input_dim]
                // Apply LSTM
                let (new_hidden, new_cell) = self.lstm_cell(&lstm_input, &hidden, &cell)?;
                hidden = new_hidden;
                cell = new_cell;
            }
            // Final representation is concatenation of memory and final hidden state
            let final_repr = Tensor::cat(&[memory, hidden], -1)?;  // [1, 2 * input_dim]
            graph_representations.push(final_repr);
        }
        if graph_representations.is_empty() {
            Tensor::zeros([0, 2 * self.input_dim], &self.device)
        } else {
            Tensor::cat(&graph_representations, 0)
        }
    }
    fn strategy(&self) -> &PoolingStrategy {
        &self.strategy
    }
    fn input_dim(&self) -> usize {
        self.input_dim
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        vec![&self.lstm_weights, &self.lstm_bias, &self.attention_projection]
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        vec![&mut self.lstm_weights, &mut self.lstm_bias, &mut self.attention_projection]
    }
 }
 /// Attention-based pooling implementation
 pub struct AttentionPooling {
    strategy: PoolingStrategy,
    input_dim: usize,
    device: Device,
    // Attention parameters
    query_projection: Tensor,
    key_projection: Tensor,
    value_projection: Tensor,
 }
 impl AttentionPooling {
    /// Create new attention pooling layer
    pub fn new(input_dim: usize, device: &Device) -> Result<Self> {
        let query_projection = Tensor::randn(&[input_dim, input_dim], device)?;
        let key_projection = Tensor::randn(&[input_dim, input_dim], device)?;
        let value_projection = Tensor::randn(&[input_dim, input_dim], device)?;
        Ok(Self {
            strategy: PoolingStrategy::Attention,
            input_dim,
            device: device.clone(),
            query_projection,
            key_projection,
            value_projection,
        })
    }
 }
 impl GraphPooling for AttentionPooling {
    fn forward(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_graphs = graph.batch_size;
        let mut graph_representations = Vec::new();
        for i in 0..num_graphs {
            let start_idx = graph.graph_boundaries[i];
            let end_idx = graph.graph_boundaries[i + 1];
            if start_idx >= end_idx {
                let empty_repr = Tensor::zeros([1, self.input_dim], &self.device)?;
                graph_representations.push(empty_repr);
                continue;
            }
            let graph_nodes = graph.node_features.narrow(0, start_idx, end_idx - start_idx)?;
            // Compute query, keys, values
            let queries = graph_nodes.matmul(&self.query_projection)?;
            let keys = graph_nodes.matmul(&self.key_projection)?;
            let values = graph_nodes.matmul(&self.value_projection)?;
            // Use mean of queries as the global query
            let global_query = queries.mean_dim(&[0], true)?;  // [1, input_dim]
            // Compute attention scores
            let scores = global_query.matmul(&keys.t()?)?;  // [1, num_nodes]
            let attention_weights = scores.softmax(-1)?;  // [1, num_nodes]
            // Apply attention to values
            let attended = attention_weights.matmul(&values)?;  // [1, input_dim]
            graph_representations.push(attended);
        }
        if graph_representations.is_empty() {
            Tensor::zeros([0, self.input_dim], &self.device)
        } else {
            Tensor::cat(&graph_representations, 0)
        }
    }
    fn strategy(&self) -> &PoolingStrategy {
        &self.strategy
    }
    fn input_dim(&self) -> usize {
        self.input_dim
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        vec![&self.query_projection, &self.key_projection, &self.value_projection]
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        vec![&mut self.query_projection, &mut self.key_projection, &mut self.value_projection]
    }
 }
 // Helper function to create appropriate pooling layer
 impl dyn GraphPooling {
    pub fn new(strategy: PoolingStrategy, input_dim: usize, device: &Device) -> Result<Box<dyn GraphPooling>> {
        match strategy {
            PoolingStrategy::Global => Ok(Box::new(GlobalPooling::new(input_dim, device)?)),
            PoolingStrategy::Hierarchical => Ok(Box::new(HierarchicalPooling::new(input_dim, 2, 0.5, device)?)),
            PoolingStrategy::Set2Set => Ok(Box::new(Set2SetPooling::new(input_dim, device)?)),
            PoolingStrategy::Attention => Ok(Box::new(AttentionPooling::new(input_dim, device)?)),
        }
    }
 }
@@ -1,483 +0,0 @@
 //! Graph Transformer implementation for attention-based learning on graph-structured data
 //!
 //! Provides complete Graph Transformer functionality with support for:
 //! - Node and edge feature processing
 //! - Multiple attention mechanisms
 //! - Graph-level representations via pooling
 //! - Batch processing of multiple graphs
 use crate::{Result, TransformerError};
 use crate::graph::{GraphBatch, Graph, GraphAttentionType};
 use rtx_tensor::{Tensor, Device, DType};
 use serde::{Deserialize, Serialize};
 /// Graph pooling strategies
 #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
 pub enum PoolingStrategy {
    /// Global mean/max/sum pooling
    Global,
    /// Hierarchical graph pooling
    Hierarchical,
    /// Set2Set pooling mechanism
    Set2Set,
    /// Attention-based pooling
    Attention,
 }
 /// Graph Transformer configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GraphTransformerConfig {
    /// Node feature dimension
    pub node_dim: usize,
    /// Edge feature dimension  
    pub edge_dim: usize,
    /// Hidden dimension for transformations
    pub hidden_dim: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Number of transformer layers
    pub num_layers: usize,
    /// Type of graph attention mechanism
    pub attention_type: GraphAttentionType,
    /// Graph pooling strategy
    pub pooling_strategy: PoolingStrategy,
    /// Dropout probability
    pub dropout: f32,
    /// Whether to use edge features
    pub use_edge_features: bool,
    /// Whether to use positional encoding
    pub use_positional_encoding: bool,
    /// Whether to return attention weights
    pub return_attention_weights: bool,
    /// Memory efficient mode
    pub memory_efficient: bool,
    /// Gradient checkpointing
    pub gradient_checkpointing: bool,
    /// Layer normalization epsilon
    pub layer_norm_eps: f64,
    /// Activation function
    pub activation: String,
 }
 impl GraphTransformerConfig {
    /// Create new Graph Transformer configuration
    pub fn new(node_dim: usize, edge_dim: usize, hidden_dim: usize, num_heads: usize) -> Self {
        if num_heads == 0 {
            panic!("Number of attention heads must be greater than 0");
        }
        if node_dim == 0 {
            panic!("Node dimension must be greater than 0");
        }
        if hidden_dim == 0 {
            panic!("Hidden dimension must be greater than 0");
        }
        if hidden_dim % num_heads != 0 {
            panic!("Hidden dimension must be divisible by number of heads");
        }
        Self {
            node_dim,
            edge_dim,
            hidden_dim,
            num_heads,
            num_layers: 6,
            attention_type: GraphAttentionType::Standard,
            pooling_strategy: PoolingStrategy::Global,
            dropout: 0.1,
            use_edge_features: true,
            use_positional_encoding: true,
            return_attention_weights: false,
            memory_efficient: false,
            gradient_checkpointing: false,
            layer_norm_eps: 1e-5,
            activation: "relu".to_string(),
        }
    }
    /// Set number of layers
    pub fn with_num_layers(mut self, num_layers: usize) -> Self {
        self.num_layers = num_layers;
        self
    }
    /// Set attention type
    pub fn with_attention_type(mut self, attention_type: GraphAttentionType) -> Self {
        self.attention_type = attention_type;
        self
    }
    /// Set pooling strategy
    pub fn with_pooling_strategy(mut self, pooling_strategy: PoolingStrategy) -> Self {
        self.pooling_strategy = pooling_strategy;
        self
    }
    /// Set dropout rate
    pub fn with_dropout(mut self, dropout: f32) -> Self {
        self.dropout = dropout;
        self
    }
    /// Disable edge features
    pub fn without_edge_features(mut self) -> Self {
        self.use_edge_features = false;
        self
    }
    /// Disable positional encoding
    pub fn without_positional_encoding(mut self) -> Self {
        self.use_positional_encoding = false;
        self
    }
    /// Enable returning attention weights
    pub fn with_return_attention_weights(mut self, return_weights: bool) -> Self {
        self.return_attention_weights = return_weights;
        self
    }
    /// Enable memory efficient mode
    pub fn with_memory_efficient(mut self, memory_efficient: bool) -> Self {
        self.memory_efficient = memory_efficient;
        self
    }
    /// Enable gradient checkpointing
    pub fn with_gradient_checkpointing(mut self, gradient_checkpointing: bool) -> Self {
        self.gradient_checkpointing = gradient_checkpointing;
        self
    }
    /// Set activation function
    pub fn with_activation(mut self, activation: &str) -> Self {
        self.activation = activation.to_string();
        self
    }
 }
 /// Output from Graph Transformer
 #[derive(Debug)]
 pub struct GraphOutput {
    /// Updated node representations [total_nodes, node_dim]
    pub node_representations: Tensor,
    /// Graph-level representations [batch_size, node_dim]  
    pub graph_representations: Tensor,
    /// Attention weights if requested [num_layers, batch_size, num_heads, max_nodes, max_nodes]
    pub attention_weights: Option<Tensor>,
    /// Updated edge representations if using edge features
    pub edge_representations: Option<Tensor>,
 }
 /// Graph Transformer implementation
 pub struct GraphTransformer {
    config: GraphTransformerConfig,
    device: Device,
    // Core components
    node_embedding: Tensor,       // Linear layer for node input projection
    edge_embedding: Option<Tensor>, // Linear layer for edge input projection
    // Transformer layers
    attention_layers: Vec<Box<dyn GraphLayer>>,
    norm_layers: Vec<Tensor>,     // Layer normalization parameters
    // Pooling layer
    pooling_layer: Box<dyn GraphPooling>,
    // Output projection
    output_projection: Tensor,
    // Positional encoding
    positional_encoding: Option<Box<dyn GraphPositionalEncoding>>,
 }
 impl GraphTransformer {
    /// Create a new Graph Transformer
    pub fn new(config: GraphTransformerConfig, device: &Device) -> Result<Self> {
        // Initialize node embedding (input projection)
        let node_embedding = Tensor::randn(&[config.node_dim, config.hidden_dim], device)?;
        // Initialize edge embedding if using edge features
        let edge_embedding = if config.use_edge_features {
            Some(Tensor::randn(&[config.edge_dim, config.hidden_dim], device)?)
        } else {
            None
        };
        // Create attention layers
        let mut attention_layers: Vec<Box<dyn GraphLayer>> = Vec::new();
        for _ in 0..config.num_layers {
            let attention_config = GraphAttentionConfig::new(
                config.hidden_dim,
                config.num_heads,
                config.attention_type,
            );
            let attention_layer = GraphAttention::new(attention_config, device)?;
            attention_layers.push(Box::new(attention_layer));
        }
        // Create layer normalization parameters
        let mut norm_layers = Vec::new();
        for _ in 0..config.num_layers {
            let norm_weight = Tensor::ones([config.hidden_dim], device)?;
            norm_layers.push(norm_weight);
        }
        // Create pooling layer
        let pooling_layer: Box<dyn GraphPooling> = match config.pooling_strategy {
            PoolingStrategy::Global => {
                Box::new(GlobalPooling::new(config.hidden_dim, device)?)
            }
            PoolingStrategy::Hierarchical => {
                Box::new(HierarchicalPooling::new(config.hidden_dim, 2, 0.5, device)?)
            }
            PoolingStrategy::Set2Set => {
                Box::new(Set2SetPooling::new(config.hidden_dim, device)?)
            }
            PoolingStrategy::Attention => {
                Box::new(AttentionPooling::new(config.hidden_dim, device)?)
            }
        };
        // Output projection
        let output_projection = Tensor::randn(&[config.hidden_dim, config.node_dim], device)?;
        // Positional encoding
        let positional_encoding = if config.use_positional_encoding {
            let pe_config = GraphPEConfig::new(config.hidden_dim, GraphPEType::Laplacian);
            Some(Box::new(GraphPositionalEncoding::new(pe_config, device)?))
        } else {
            None
        };
        Ok(Self {
            config,
            device: device.clone(),
            node_embedding,
            edge_embedding,
            attention_layers,
            norm_layers,
            pooling_layer,
            output_projection,
            positional_encoding,
        })
    }
    /// Get the configuration
    pub fn config(&self) -> &GraphTransformerConfig {
        &self.config
    }
 }
 impl GraphLayer for GraphTransformer {
    fn forward(&self, graph: &GraphBatch) -> Result<GraphOutput> {
        // Input validation
        if graph.num_nodes() == 0 {
            return Err(TransformerError::InvalidInput(
                "Cannot process empty graph batch".to_string()
            ));
        }
        let expected_node_dim = self.config.node_dim;
        let actual_node_dim = graph.node_features.shape().dims()[1];
        if actual_node_dim != expected_node_dim {
            return Err(TransformerError::shape_mismatch(format!(
                "Expected node features dimension {}, got {}",
                expected_node_dim, actual_node_dim
            )));
        }
        // Project input node features to hidden dimension
        let mut node_features = graph.node_features.matmul(&self.node_embedding)?;
        // Add positional encoding if enabled
        if let Some(ref pe) = self.positional_encoding {
            let pos_encoding = pe.forward(graph)?;
            node_features = node_features.add(&pos_encoding)?;
        }
        // Process edge features if enabled
        let mut edge_features = if self.config.use_edge_features && self.edge_embedding.is_some() {
            Some(graph.edge_features.matmul(self.edge_embedding.as_ref().unwrap())?)
        } else {
            None
        };
        // Storage for attention weights if requested
        let mut all_attention_weights = Vec::new();
        // Apply transformer layers
        for (layer_idx, attention_layer) in self.attention_layers.iter().enumerate() {
            // Create temporary graph batch with current features
            let temp_graph = GraphBatch::new(
                node_features.clone(),
                edge_features.clone().unwrap_or_else(|| graph.edge_features.clone()),
                graph.edge_indices.clone(),
                graph.graph_boundaries.clone(),
                self.device.clone(),
            )?;
            // Apply attention layer
            let layer_output = attention_layer.forward(&temp_graph)?;
            // Residual connection and layer norm
            let residual = node_features.clone();
            node_features = layer_output.node_representations;
            node_features = node_features.add(&residual)?;
            // Layer normalization
            node_features = self.apply_layer_norm(&node_features, layer_idx)?;
            // Store attention weights if requested
            if self.config.return_attention_weights {
                if let Some(attention_weights) = layer_output.attention_weights {
                    all_attention_weights.push(attention_weights);
                }
            }
            // Update edge features if available
            if let Some(edge_repr) = layer_output.edge_representations {
                edge_features = Some(edge_repr);
            }
        }
        // Apply pooling to get graph-level representations
        let temp_graph_for_pooling = GraphBatch::new(
            node_features.clone(),
            edge_features.clone().unwrap_or_else(|| graph.edge_features.clone()),
            graph.edge_indices.clone(),
            graph.graph_boundaries.clone(),
            self.device.clone(),
        )?;
        let graph_representations = self.pooling_layer.forward(&temp_graph_for_pooling)?;
        // Project back to original node dimension
        let final_node_representations = node_features.matmul(&self.output_projection)?;
        // Combine attention weights if requested
        let combined_attention_weights = if self.config.return_attention_weights && !all_attention_weights.is_empty() {
            Some(self.combine_attention_weights(all_attention_weights, graph.batch_size)?)
        } else {
            None
        };
        Ok(GraphOutput {
            node_representations: final_node_representations,
            graph_representations,
            attention_weights: combined_attention_weights,
            edge_representations: edge_features,
        })
    }
    fn layer_type(&self) -> &'static str {
        "GraphTransformer"
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        let mut params = vec![&self.node_embedding, &self.output_projection];
        if let Some(ref edge_emb) = self.edge_embedding {
            params.push(edge_emb);
        }
        for norm_layer in &self.norm_layers {
            params.push(norm_layer);
        }
        // Add attention layer parameters
        for attention_layer in &self.attention_layers {
            params.extend(attention_layer.parameters());
        }
        // Add pooling layer parameters
        params.extend(self.pooling_layer.parameters());
        params
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![&mut self.node_embedding, &mut self.output_projection];
        if let Some(ref mut edge_emb) = self.edge_embedding {
            params.push(edge_emb);
        }
        for norm_layer in &mut self.norm_layers {
            params.push(norm_layer);
        }
        // Add attention layer parameters
        for attention_layer in &mut self.attention_layers {
            params.extend(attention_layer.parameters_mut());
        }
        // Add pooling layer parameters
        params.extend(self.pooling_layer.parameters_mut());
        params
    }
 }
 impl GraphTransformer {
    /// Apply layer normalization
    fn apply_layer_norm(&self, input: &Tensor, layer_idx: usize) -> Result<Tensor> {
        // Simple layer normalization implementation
        let mean = input.mean_dim(&[-1], true)?;
        let variance = input.var_dim(&[-1], true, true)?;
        let normalized = input.sub(&mean)?.div(&(variance.add(&Tensor::full(&[], (self.config.layer_norm_eps) as f32, &&self.device)?))?.sqrt()?)?;
        let weight = &self.norm_layers[layer_idx];
        normalized.mul(weight)
    }
    /// Combine attention weights from all layers
    fn combine_attention_weights(&self, weights: Vec<Tensor>, batch_size: usize) -> Result<Tensor> {
        // Stack attention weights from all layers
        // Expected shape: [num_layers, batch_size, num_heads, max_nodes, max_nodes]
        if weights.is_empty() {
            return Err(TransformerError::InvalidInput(
                "No attention weights to combine".to_string()
            ));
        }
        // For now, return the last layer's attention weights
        // In a full implementation, we would stack all layers
        Ok(weights.into_iter().last().unwrap())
    }
 }
 impl Layer for GraphTransformer {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        Err(TransformerError::InvalidInput(
            "GraphTransformer requires GraphBatch input, not Tensor. Use GraphLayer::forward instead.".to_string()
        ))
    }
    fn layer_type(&self) -> &'static str {
        "GraphTransformer"
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        GraphLayer::parameters(self)
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        GraphLayer::parameters_mut(self)
    }
 }
 // Forward declarations for components that will be implemented in other files
 use crate::graph::graph_attention::{GraphAttention, GraphAttentionConfig};
 use crate::graph::graph_pooling::{GraphPooling, GlobalPooling, HierarchicalPooling, Set2SetPooling, AttentionPooling};
 use crate::graph::positional_encoding::{GraphPositionalEncoding, GraphPEConfig, GraphPEType};
@@ -1,345 +0,0 @@
 //! Graph Positional Encoding implementations
 //!
 //! Provides multiple strategies for encoding positional information in graphs:
 //! - Laplacian eigenvector-based encoding  
 //! - Random walk-based encoding
 //! - Learned positional encoding
 //! - Distance-based encoding
 use crate::{Result, TransformerError};
 use crate::graph::GraphBatch;
 use rtx_tensor::{Tensor, Device};
 use serde::{Deserialize, Serialize};
 /// Graph positional encoding types
 #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
 pub enum GraphPEType {
    /// Laplacian eigenvector-based encoding
    Laplacian,
    /// Random walk-based encoding
    RandomWalk, 
    /// Learned positional encoding
    Learned,
    /// Shortest path distance-based encoding
    Distance,
 }
 /// Graph positional encoding configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GraphPEConfig {
    /// Embedding dimension
    pub embedding_dim: usize,
    /// Type of positional encoding
    pub pe_type: GraphPEType,
    /// Number of eigenvectors to use (for Laplacian)
    pub num_eigenvectors: usize,
    /// Random walk length (for RandomWalk)
    pub walk_length: usize,
    /// Maximum number of nodes (for Learned)
    pub max_nodes: usize,
    /// Maximum distance (for Distance-based)
    pub max_distance: usize,
    /// Whether to normalize encodings
    pub normalize: bool,
 }
 impl GraphPEConfig {
    /// Create new graph PE configuration
    pub fn new(embedding_dim: usize, pe_type: GraphPEType) -> Self {
        Self {
            embedding_dim,
            pe_type,
            num_eigenvectors: 16,
            walk_length: 10,
            max_nodes: 1000,
            max_distance: 10,
            normalize: true,
        }
    }
    /// Set number of eigenvectors
    pub fn with_num_eigenvectors(mut self, num_eigenvectors: usize) -> Self {
        self.num_eigenvectors = num_eigenvectors;
        self
    }
    /// Set random walk length
    pub fn with_walk_length(mut self, walk_length: usize) -> Self {
        self.walk_length = walk_length;
        self
    }
    /// Set maximum nodes
    pub fn with_max_nodes(mut self, max_nodes: usize) -> Self {
        self.max_nodes = max_nodes;
        self
    }
 }
 /// Graph positional encoding implementation
 pub struct GraphPositionalEncoding {
    config: GraphPEConfig,
    device: Device,
    // Encoding parameters
    encoding_weights: Option<Tensor>,      // For learned encodings
    projection_layer: Tensor,              // Project PE to embedding_dim
    // Precomputed encodings (for efficiency)
    cached_encodings: Option<Tensor>,
 }
 impl GraphPositionalEncoding {
    /// Create new graph positional encoding
    pub fn new(config: GraphPEConfig, device: &Device) -> Result<Self> {
        // Initialize projection layer based on PE type
        let projection_input_dim = match config.pe_type {
            GraphPEType::Laplacian => config.num_eigenvectors,
            GraphPEType::RandomWalk => config.walk_length,
            GraphPEType::Learned => config.embedding_dim,
            GraphPEType::Distance => config.max_distance,
        };
        let projection_layer = Tensor::randn(&[projection_input_dim, config.embedding_dim], device)?;
        // Initialize learned encoding weights if needed
        let encoding_weights = if config.pe_type == GraphPEType::Learned {
            Some(Tensor::randn(&[config.max_nodes, config.embedding_dim], device)?)
        } else {
            None
        };
        Ok(Self {
            config,
            device: device.clone(),
            encoding_weights,
            projection_layer,
            cached_encodings: None,
        })
    }
    /// Get the configuration
    pub fn config(&self) -> &GraphPEConfig {
        &self.config
    }
    /// Forward pass - compute positional encodings for graph batch
    pub fn forward(&self, graph: &GraphBatch) -> Result<Tensor> {
        match self.config.pe_type {
            GraphPEType::Laplacian => self.laplacian_encoding(graph),
            GraphPEType::RandomWalk => self.random_walk_encoding(graph),
            GraphPEType::Learned => self.learned_encoding(graph),
            GraphPEType::Distance => self.distance_encoding(graph),
        }
    }
    /// Compute Laplacian eigenvector-based positional encoding
    fn laplacian_encoding(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        // Construct graph Laplacian matrix
        let laplacian = self.compute_graph_laplacian(graph)?;  // [num_nodes, num_nodes]
        // Compute eigenvectors (simplified - in practice would use proper eigendecomposition)
        // For now, we'll create a placeholder that has the right shape
        let eigenvalues = Tensor::randn(&[num_nodes, self.config.num_eigenvectors], &self.device)?;
        // Select top-k eigenvectors (excluding the first constant eigenvector)
        let pe_features = if self.config.num_eigenvectors < num_nodes {
            eigenvalues.narrow(1, 1, self.config.num_eigenvectors)?  // Skip first eigenvector
        } else {
            eigenvalues.clone()
        };
        // Project to embedding dimension
        let projected = pe_features.matmul(&self.projection_layer)?;
        // Normalize if requested
        if self.config.normalize {
            let norm = projected.norm_dim(&[1], true, true)?;
            projected.div(&(norm.add(&Tensor::full(&[], (1e-8) as f32, &&self.device)?))?)?
        } else {
            projected
        }
    }
    /// Compute random walk-based positional encoding
    fn random_walk_encoding(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        // Construct adjacency matrix
        let adjacency = self.compute_adjacency_matrix(graph)?;  // [num_nodes, num_nodes]
        // Compute transition matrix (row-normalized adjacency)
        let degree = adjacency.sum_dim(&[1], true)?;  // [num_nodes, 1]
        let degree_inv = degree.reciprocal()?;
        let transition_matrix = adjacency.mul(&degree_inv)?;
        // Compute random walk probabilities for different steps
        let mut walk_features = Vec::new();
        let mut current_probs = Tensor::eye(num_nodes, &self.device)?;  // Identity matrix
        for _ in 0..self.config.walk_length {
            current_probs = current_probs.matmul(&transition_matrix)?;
            // Use diagonal values as features (probability of returning to starting node)
            let diag = current_probs.diag()?;  // [num_nodes]
            walk_features.push(diag);
        }
        // Stack walk features
        let stacked_features = Tensor::stack(&walk_features, 1)?;  // [num_nodes, walk_length]
        // Project to embedding dimension
        let projected = stacked_features.matmul(&self.projection_layer)?;
        // Normalize if requested
        if self.config.normalize {
            let norm = projected.norm_dim(&[1], true, true)?;
            projected.div(&(norm.add(&Tensor::full(&[], (1e-8) as f32, &&self.device)?))?)?
        } else {
            projected
        }
    }
    /// Compute learned positional encoding
    fn learned_encoding(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        if let Some(ref weights) = self.encoding_weights {
            // Simple learned encoding based on node indices
            let node_indices = Tensor::arange(0, num_nodes as i64, &self.device)?;
            let encodings = weights.index_select(0, &node_indices)?;
            if self.config.normalize {
                let norm = encodings.norm_dim(&[1], true, true)?;
                encodings.div(&(norm.add(&Tensor::full(&[], (1e-8) as f32, &&self.device)?))?)?
            } else {
                encodings
            }
        } else {
            Err(TransformerError::InvalidInput(
                "Encoding weights not initialized for learned PE".to_string()
            ))
        }
    }
    /// Compute distance-based positional encoding
    fn distance_encoding(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        // Compute shortest path distances (simplified implementation)
        let distance_matrix = self.compute_shortest_paths(graph)?;  // [num_nodes, num_nodes]
        // Create distance-based features
        // For each node, create a histogram of distances to other nodes
        let mut distance_features = Vec::new();
        for dist in 1..=self.config.max_distance {
            let dist_tensor = Tensor::full(&[], (dist as f32) as f32, &&self.device)?;
            let mask = distance_matrix.eq(&dist_tensor)?;
            let count = mask.sum_dim(&[1], false)?.to_dtype(rtx_tensor::DType::F32)?;
            distance_features.push(count);
        }
        // Stack distance features
        let stacked_features = Tensor::stack(&distance_features, 1)?;  // [num_nodes, max_distance]
        // Project to embedding dimension
        let projected = stacked_features.matmul(&self.projection_layer)?;
        // Normalize if requested
        if self.config.normalize {
            let norm = projected.norm_dim(&[1], true, true)?;
            projected.div(&(norm.add(&Tensor::full(&[], (1e-8) as f32, &&self.device)?))?)?
        } else {
            projected
        }
    }
    /// Compute graph Laplacian matrix
    fn compute_graph_laplacian(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        // Construct adjacency matrix
        let adjacency = self.compute_adjacency_matrix(graph)?;
        // Compute degree matrix
        let degree_values = adjacency.sum_dim(&[1], false)?;  // [num_nodes]
        let degree_matrix = Tensor::diag(&degree_values)?;    // [num_nodes, num_nodes]
        // Laplacian = D - A
        degree_matrix.sub(&adjacency)
    }
    /// Compute adjacency matrix from edge indices
    fn compute_adjacency_matrix(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        let num_edges = graph.num_edges();
        // Initialize adjacency matrix
        let mut adjacency = Tensor::zeros([num_nodes, num_nodes], &self.device)?;
        // Fill adjacency matrix based on edges
        // This is a simplified implementation - in practice would use scatter operations
        let source_indices = graph.edge_indices.get(&[0])?;  // [num_edges]
        let target_indices = graph.edge_indices.get(&[1])?;  // [num_edges]
        // For now, create a simple symmetric adjacency matrix
        for i in 0..num_edges {
            // In a real implementation, we would set adjacency[source[i], target[i]] = 1
            // This is a placeholder implementation
        }
        // Return identity + small random values as placeholder
        let identity = Tensor::eye(num_nodes, &self.device)?;
        let noise = Tensor::randn(&[num_nodes, num_nodes], &self.device)?.mul(&Tensor::full(&[], 0.1, &self.device))?;
        identity.add(&noise)
    }
    /// Compute shortest path distances (simplified Floyd-Warshall-like)
    fn compute_shortest_paths(&self, graph: &GraphBatch) -> Result<Tensor> {
        let num_nodes = graph.num_nodes();
        // Initialize distance matrix with large values
        let mut distances = Tensor::full([num_nodes, num_nodes], f32::INFINITY, &self.device)?;
        // Set diagonal to 0
        let diagonal_indices = Tensor::arange(0, num_nodes as i64, &self.device)?;
        // distances[diagonal_indices, diagonal_indices] = 0  (simplified)
        // Set direct edges to distance 1
        // This would require scatter operations in a full implementation
        // For now, return a simple distance matrix as placeholder
        let identity = Tensor::eye(num_nodes, &self.device)?;
        let ones = Tensor::ones([num_nodes, num_nodes], &self.device)?;
        identity.add(&ones)  // Distance matrix with 1s for neighbors, 2s otherwise
    }
    /// Get parameters for optimization
    pub fn parameters(&self) -> Vec<&Tensor> {
        let mut params = vec![&self.projection_layer];
        if let Some(ref encoding_weights) = self.encoding_weights {
            params.push(encoding_weights);
        }
        params
    }
    /// Get mutable parameters for optimization
    pub fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![&mut self.projection_layer];
        if let Some(ref mut encoding_weights) = self.encoding_weights {
            params.push(encoding_weights);
        }
        params
    }
    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }
 }
@@ -48,8 +48,6 @@ pub mod sparse_attention;
 // MoE (Mixture of Experts)
 pub mod metal_moe;
 pub mod mixture_of_experts;
 // pub mod moe_layer;
 // pub mod moe_integration;
 // pub mod moe_routing;
 // pub mod capacity_tuning;
 // pub mod expert_dropout;
@@ -1,564 +0,0 @@
 //! Integration layer for MoE with existing transformer architectures
 use crate::{Result, TransformerError};
 use crate::layers::{Layer, MoEConfig, MoELayer, MoEOutput};
 use crate::architectures::TransformerConfig;
 use rtx_tensor::{Tensor, Device, DType};
 use rtx_autograd::TensorAutograd;
 use serde::{Deserialize, Serialize};
 /// Configuration for MoE-enabled transformer layers
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MoETransformerConfig {
    /// Base transformer configuration
    pub transformer_config: TransformerConfig,
    /// MoE configuration (replaces traditional FFN)
    pub moe_config: MoEConfig,
    /// Whether to use MoE in all layers or only specific ones
    pub moe_layers: Vec<usize>, // Layer indices that should use MoE
    /// Whether to keep traditional FFN alongside MoE (hybrid mode)
    pub hybrid_mode: bool,
    /// Weight for combining traditional FFN with MoE in hybrid mode
    pub hybrid_weight: f32,
 }
 impl Default for MoETransformerConfig {
    fn default() -> Self {
        Self {
            transformer_config: TransformerConfig::default(),
            moe_config: MoEConfig::default(),
            moe_layers: vec![],
            hybrid_mode: false,
            hybrid_weight: 0.5,
        }
    }
 }
 impl MoETransformerConfig {
    /// Create a new MoE transformer configuration
    pub fn new(transformer_config: TransformerConfig, moe_config: MoEConfig) -> Self {
        Self {
            transformer_config,
            moe_config,
            ..Default::default()
        }
    }
    /// Enable MoE in all transformer layers
    pub fn enable_moe_all_layers(&mut self) {
        self.moe_layers = (0..self.transformer_config.num_layers).collect();
    }
    /// Enable MoE in specific layers
    pub fn enable_moe_layers(&mut self, layer_indices: Vec<usize>) -> Result<()> {
        for &idx in &layer_indices {
            if idx >= self.transformer_config.num_layers {
                return Err(TransformerError::config(
                    format!("MoE layer index {} exceeds total layers {}", 
                           idx, self.transformer_config.num_layers)
                ));
            }
        }
        self.moe_layers = layer_indices;
        Ok(())
    }
    /// Enable hybrid mode with traditional FFN + MoE
    pub fn enable_hybrid_mode(&mut self, hybrid_weight: f32) -> Result<()> {
        if !(0.0..=1.0).contains(&hybrid_weight) {
            return Err(TransformerError::config(
                "hybrid_weight must be between 0.0 and 1.0".to_string()
            ));
        }
        self.hybrid_mode = true;
        self.hybrid_weight = hybrid_weight;
        Ok(())
    }
    /// Check if a given layer should use MoE
    pub fn is_moe_layer(&self, layer_index: usize) -> bool {
        self.moe_layers.contains(&layer_index)
    }
    /// Validate the configuration
    pub fn validate(&self) -> Result<()> {
        self.transformer_config.validate()?;
        self.moe_config.validate()?;
        // Ensure MoE hidden dimensions match transformer
        if self.moe_config.hidden_dim != self.transformer_config.d_model {
            return Err(TransformerError::config(
                format!("MoE hidden_dim ({}) must match transformer d_model ({})", 
                       self.moe_config.hidden_dim, self.transformer_config.d_model)
            ));
        }
        // Check layer indices
        for &idx in &self.moe_layers {
            if idx >= self.transformer_config.num_layers {
                return Err(TransformerError::config(
                    format!("MoE layer index {} exceeds total layers {}", 
                           idx, self.transformer_config.num_layers)
                ));
            }
        }
        // Validate hybrid mode settings
        if self.hybrid_mode && !(0.0..=1.0).contains(&self.hybrid_weight) {
            return Err(TransformerError::config(
                "hybrid_weight must be between 0.0 and 1.0".to_string()
            ));
        }
        Ok(())
    }
    /// Estimate total memory usage
    pub fn estimate_memory_usage(&self, batch_size: usize, seq_len: usize) -> usize {
        let base_memory = self.transformer_config.estimate_kv_cache_memory(seq_len);
        let moe_memory = self.moe_config.estimate_memory_usage(batch_size, seq_len);
        // Memory for MoE layers
        let moe_layer_memory = moe_memory * self.moe_layers.len();
        base_memory + moe_layer_memory
    }
 }
 /// MoE-enabled feedforward layer that can replace traditional FFN
 #[derive(Debug)]
 pub struct MoEFeedForward {
    /// Configuration
    config: MoETransformerConfig,
    /// Device
    device: Device,
    /// MoE layer
    moe_layer: MoELayer,
    /// Traditional FFN for hybrid mode (optional)
    traditional_ffn: Option<TraditionalFFN>,
 }
 impl MoEFeedForward {
    /// Create a new MoE feedforward layer
    pub fn new(config: MoETransformerConfig, device: &Device) -> Result<Self> {
        config.validate()?;
        let moe_layer = MoELayer::new(config.moe_config.clone(), device)?;
        let traditional_ffn = if config.hybrid_mode {
            Some(TraditionalFFN::new(
                config.transformer_config.d_model,
                config.transformer_config.d_ff,
                &config.transformer_config.activation,
                config.transformer_config.bias,
                device
            )?)
        } else {
            None
        };
        Ok(Self {
            config,
            device: device.clone(),
            moe_layer,
            traditional_ffn,
        })
    }
    /// Forward pass with optional auxiliary loss output
    pub fn forward_with_aux(&self, input: &Tensor) -> Result<(Tensor, Option<Tensor>)> {
        let moe_output = self.moe_layer.forward(input)?;
        let output = if let Some(ref traditional_ffn) = self.traditional_ffn {
            // Hybrid mode: combine MoE and traditional FFN
            let ffn_output = traditional_ffn.forward(input)?;
            let moe_weighted = moe_output.output.mul_scalar(self.config.hybrid_weight)?;
            let ffn_weighted = ffn_output.mul_scalar(1.0 - self.config.hybrid_weight)?;
            moe_weighted.add(&ffn_weighted)?
        } else {
            // Pure MoE mode
            moe_output.output
        };
        let aux_loss = moe_output.aux_loss().cloned();
        Ok((output, aux_loss))
    }
    /// Get MoE output with full statistics
    pub fn forward_detailed(&self, input: &Tensor) -> Result<MoEDetailedOutput> {
        let moe_output = self.moe_layer.forward(input)?;
        let final_output = if let Some(ref traditional_ffn) = self.traditional_ffn {
            // Hybrid mode
            let ffn_output = traditional_ffn.forward(input)?;
            let moe_weighted = moe_output.output.mul_scalar(self.config.hybrid_weight)?;
            let ffn_weighted = ffn_output.mul_scalar(1.0 - self.config.hybrid_weight)?;
            let combined = moe_weighted.add(&ffn_weighted)?;
            Some(MoEDetailedOutput {
                output: combined,
                moe_output: moe_output.output,
                ffn_output: Some(ffn_output),
                aux_loss: moe_output.aux_loss().cloned(),
                routing_info: moe_output.routing_info,
                load_stats: moe_output.load_stats,
                is_load_balanced: moe_output.is_load_balanced(),
            })
        } else {
            // Pure MoE mode
            Some(MoEDetailedOutput {
                output: moe_output.output.clone(),
                moe_output: moe_output.output,
                ffn_output: None,
                aux_loss: moe_output.aux_loss().cloned(),
                routing_info: moe_output.routing_info,
                load_stats: moe_output.load_stats,
                is_load_balanced: moe_output.is_load_balanced(),
            })
        };
        final_output.ok_or_else(|| TransformerError::runtime("Failed to create detailed output".to_string()))
    }
    /// Get all parameters including MoE and optional traditional FFN
    pub fn parameters(&self) -> Vec<&Tensor> {
        let mut params = self.moe_layer.parameters();
        if let Some(ref ffn) = self.traditional_ffn {
            params.extend(ffn.parameters());
        }
        params
    }
    /// Get all mutable parameters
    pub fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = self.moe_layer.parameters_mut();
        if let Some(ref mut ffn) = self.traditional_ffn {
            params.extend(ffn.parameters_mut());
        }
        params
    }
    /// Get configuration
    pub fn config(&self) -> &MoETransformerConfig {
        &self.config
    }
 }
 impl Layer for MoEFeedForward {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let (output, _aux_loss) = self.forward_with_aux(input)?;
        Ok(output)
    }
    fn layer_type(&self) -> &'static str {
        if self.traditional_ffn.is_some() {
            "HybridMoEFeedForward"
        } else {
            "MoEFeedForward"
        }
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        self.parameters()
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        self.parameters_mut()
    }
 }
 /// Traditional feedforward network for hybrid mode
 #[derive(Debug)]
 struct TraditionalFFN {
    up_weight: Tensor,
    up_bias: Option<Tensor>,
    down_weight: Tensor,
    down_bias: Option<Tensor>,
    activation: String,
    device: Device,
 }
 impl TraditionalFFN {
    fn new(
        hidden_dim: usize,
        intermediate_dim: usize,
        activation: &str,
        use_bias: bool,
        device: &Device,
    ) -> Result<Self> {
        let up_weight = Tensor::randn(&[intermediate_dim, hidden_dim], DType::F32, device)?
            .require_grad()?;
        let up_bias = if use_bias {
            Some(Tensor::zeros(&[intermediate_dim], device)?.require_grad())
        } else {
            None
        };
        let down_weight = Tensor::randn(&[hidden_dim, intermediate_dim], DType::F32, device)?
            .require_grad()?;
        let down_bias = if use_bias {
            Some(Tensor::zeros(&[hidden_dim], device)?.require_grad())
        } else {
            None
        };
        Ok(Self {
            up_weight,
            up_bias,
            down_weight,
            down_bias,
            activation: activation.to_string(),
            device: device.clone(),
        })
    }
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        // Up projection
        let up_output = input.matmul(&self.up_weight.t()?)?;
        let up_output = if let Some(ref bias) = self.up_bias {
            up_output.add(bias)?
        } else {
            up_output
        };
        // Activation
        let activated = self.apply_activation(&up_output)?;
        // Down projection
        let down_output = activated.matmul(&self.down_weight.t()?)?;
        let output = if let Some(ref bias) = self.down_bias {
            down_output.add(bias)?
        } else {
            down_output
        };
        Ok(output)
    }
    fn apply_activation(&self, input: &Tensor) -> Result<Tensor> {
        match self.activation.as_str() {
            "relu" => input.relu(),
            "gelu" => input.gelu(),
            "swish" | "silu" => input.swish(),
            "tanh" => input.tanh(),
            _ => Err(TransformerError::config(
                format!("Unsupported activation function: {}", self.activation)
            )),
        }
    }
    fn parameters(&self) -> Vec<&Tensor> {
        let mut params = vec![&self.up_weight, &self.down_weight];
        if let Some(ref bias) = self.up_bias {
            params.push(bias);
        }
        if let Some(ref bias) = self.down_bias {
            params.push(bias);
        }
        params
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = vec![&mut self.up_weight, &mut self.down_weight];
        if let Some(ref mut bias) = self.up_bias {
            params.push(bias);
        }
        if let Some(ref mut bias) = self.down_bias {
            params.push(bias);
        }
        params
    }
 }
 /// Detailed output from MoE feedforward layer
 #[derive(Debug)]
 pub struct MoEDetailedOutput {
    /// Final combined output
    pub output: Tensor,
    /// MoE component output
    pub moe_output: Tensor,
    /// Traditional FFN output (in hybrid mode)
    pub ffn_output: Option<Tensor>,
    /// Auxiliary loss for training
    pub aux_loss: Option<Tensor>,
    /// Routing information
    pub routing_info: crate::layers::RoutingInfo,
    /// Load balancing statistics
    pub load_stats: crate::layers::LoadBalancingStats,
    /// Whether the load is balanced
    pub is_load_balanced: bool,
 }
 // Implementation for TransformerConfig validation
 impl TransformerConfig {
    /// Validate transformer configuration
    pub fn validate(&self) -> Result<()> {
        if self.vocab_size == 0 {
            return Err(TransformerError::config("vocab_size must be greater than 0".to_string()));
        }
        if self.d_model == 0 {
            return Err(TransformerError::config("d_model must be greater than 0".to_string()));
        }
        if self.num_layers == 0 {
            return Err(TransformerError::config("num_layers must be greater than 0".to_string()));
        }
        if self.num_heads == 0 {
            return Err(TransformerError::config("num_heads must be greater than 0".to_string()));
        }
        if self.d_model % self.num_heads != 0 {
            return Err(TransformerError::config(
                format!("d_model ({}) must be divisible by num_heads ({})", 
                       self.d_model, self.num_heads)
            ));
        }
        Ok(())
    }
 }
 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
    use super::*;
    #[test]
    fn test_moe_transformer_config() {
        let transformer_config = TransformerConfig::new(10000, 768, 12, 12, 3072, 1024);
        let moe_config = MoEConfig::new(8, 2, 768, 3072);
        let mut config = MoETransformerConfig::new(transformer_config, moe_config);
        assert!(config.validate().is_ok());
        // Test enabling MoE in all layers
        config.enable_moe_all_layers();
        assert_eq!(config.moe_layers.len(), 12);
        // Test enabling MoE in specific layers
        config.enable_moe_layers(vec![0, 2, 4, 6]).unwrap();
        assert_eq!(config.moe_layers, vec![0, 2, 4, 6]);
        // Test hybrid mode
        config.enable_hybrid_mode(0.7).unwrap();
        assert!(config.hybrid_mode);
        assert_eq!(config.hybrid_weight, 0.7);
    }
    #[test]
    fn test_moe_transformer_config_validation() {
        let transformer_config = TransformerConfig::new(10000, 768, 12, 12, 3072, 1024);
        let mut moe_config = MoEConfig::new(8, 2, 512, 3072); // Wrong hidden_dim
        let config = MoETransformerConfig::new(transformer_config, moe_config);
        // Should fail due to mismatched dimensions
        assert!(config.validate().is_err());
        // Fix dimensions
        moe_config.hidden_dim = 768;
        let mut config = MoETransformerConfig::new(transformer_config, moe_config);
        // Should pass now
        assert!(config.validate().is_ok());
        // Test invalid layer indices
        assert!(config.enable_moe_layers(vec![15]).is_err()); // Layer 15 doesn't exist
        // Test invalid hybrid weight
        assert!(config.enable_hybrid_mode(1.5).is_err());
        assert!(config.enable_hybrid_mode(-0.1).is_err());
    }
    #[test]
    fn test_moe_feedforward_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let transformer_config = TransformerConfig::new(10000, 512, 6, 8, 2048, 512);
        let moe_config = MoEConfig::new(4, 2, 512, 2048);
        let config = MoETransformerConfig::new(transformer_config, moe_config);
        let moe_ff = MoEFeedForward::new(config, &device);
        assert!(moe_ff.is_ok());
        let moe_ff = moe_ff.unwrap();
        assert_eq!(moe_ff.layer_type(), "MoEFeedForward");
    }
    #[test]
    fn test_moe_feedforward_hybrid_mode() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let transformer_config = TransformerConfig::new(10000, 512, 6, 8, 2048, 512);
        let moe_config = MoEConfig::new(4, 2, 512, 2048);
        let mut config = MoETransformerConfig::new(transformer_config, moe_config);
        config.enable_hybrid_mode(0.6).unwrap();
        let moe_ff = MoEFeedForward::new(config, &device).unwrap();
        assert_eq!(moe_ff.layer_type(), "HybridMoEFeedForward");
        // Test forward pass
        let input = Tensor::randn(&[2, 8, 512], DType::F32, &device).unwrap();
        let (output, aux_loss) = moe_ff.forward_with_aux(&input).unwrap();
        // Output shape should match input
        assert_eq!(output.shape(), input.shape());
        // Should have auxiliary loss
        assert!(aux_loss.is_some());
    }
    #[test]
    fn test_moe_feedforward_detailed_output() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let transformer_config = TransformerConfig::new(10000, 256, 4, 4, 1024, 256);
        let moe_config = MoEConfig::new(4, 2, 256, 1024);
        let mut config = MoETransformerConfig::new(transformer_config, moe_config);
        config.enable_hybrid_mode(0.5).unwrap();
        let moe_ff = MoEFeedForward::new(config, &device).unwrap();
        let input = Tensor::randn(&[1, 4, 256], DType::F32, &device).unwrap();
        let detailed_output = moe_ff.forward_detailed(&input).unwrap();
        // Should have both MoE and FFN outputs in hybrid mode
        assert!(detailed_output.ffn_output.is_some());
        // Outputs should have correct shapes
        assert_eq!(detailed_output.output.shape(), input.shape());
        assert_eq!(detailed_output.moe_output.shape(), input.shape());
        assert_eq!(detailed_output.ffn_output.as_ref().unwrap().shape(), input.shape());
        // Should have routing information
        assert_eq!(detailed_output.routing_info.expert_token_counts.len(), 4);
    }
    #[test]
    fn test_traditional_ffn() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let ffn = TraditionalFFN::new(256, 1024, "gelu", true, &device).unwrap();
        let input = Tensor::randn(&[2, 8, 256], DType::F32, &device).unwrap();
        let output = ffn.forward(&input).unwrap();
        // Output shape should match input
        assert_eq!(output.shape(), input.shape());
        // Should have correct number of parameters (with bias)
        let params = ffn.parameters();
        assert_eq!(params.len(), 4); // up_weight, down_weight, up_bias, down_bias
    }
    #[test]
    fn test_memory_estimation() {
        let transformer_config = TransformerConfig::new(10000, 768, 12, 12, 3072, 1024);
        let moe_config = MoEConfig::new(8, 2, 768, 3072);
        let mut config = MoETransformerConfig::new(transformer_config, moe_config);
        config.enable_moe_layers(vec![0, 2, 4, 6]).unwrap();
        let memory_usage = config.estimate_memory_usage(32, 128);
        assert!(memory_usage > 0);
        // More MoE layers should use more memory
        config.enable_moe_all_layers();
        let full_memory_usage = config.estimate_memory_usage(32, 128);
        assert!(full_memory_usage > memory_usage);
    }
 }
@@ -1,497 +0,0 @@
 //! MoE Layer implementation with load balancing and auxiliary losses
 use crate::layers::Layer;
 use crate::{Result, TransformerError};
 use crate::layers::mixture_of_experts::{MoEConfig, Router, Expert, RoutingInfo};
 use rtx_tensor::{Tensor, Device, DType};
 use rtx_autograd::TensorAutograd;
 use serde::{Deserialize, Serialize};
 /// Load balancer for ensuring even distribution of tokens across experts
 #[derive(Debug)]
 pub struct LoadBalancer {
    /// Configuration
    config: MoEConfig,
    /// Device
    device: Device,
 }
 impl LoadBalancer {
    /// Create a new load balancer
    pub fn new(config: MoEConfig, device: &Device) -> Result<Self> {
        config.validate()?;
        Ok(Self {
            config,
            device: device.clone(),
        })
    }
    /// Apply capacity-based load balancing to routing decisions
    pub fn apply_capacity_constraints(&self, routing_info: &RoutingInfo, capacity_per_expert: usize) -> Result<RoutingInfo> {
        // This is a simplified implementation
        // In practice, we would need to:
        // 1. Sort tokens by routing weights for each expert
        // 2. Keep only top tokens up to capacity limit
        // 3. Reassign overflow tokens to other experts or drop them
        // For now, return the input routing info unchanged
        Ok(routing_info.clone())
    }
    /// Calculate load balancing statistics
    pub fn calculate_load_stats(&self, expert_counts: &[usize], total_tokens: usize) -> LoadBalancingStats {
        if expert_counts.is_empty() || total_tokens == 0 {
            return LoadBalancingStats::default();
        }
        let average_load = total_tokens as f32 / self.config.num_experts as f32;
        let mut imbalance = 0.0;
        let mut min_load = expert_counts[0] as f32;
        let mut max_load = expert_counts[0] as f32;
        for &count in expert_counts {
            let count_f32 = count as f32;
            min_load = min_load.min(count_f32);
            max_load = max_load.max(count_f32);
            imbalance += (count_f32 - average_load).abs();
        }
        LoadBalancingStats {
            average_load,
            min_load,
            max_load,
            imbalance_score: imbalance / (self.config.num_experts as f32 * average_load).max(1e-6),
            expert_utilization: expert_counts.iter().map(|&c| c > 0).count() as f32 / self.config.num_experts as f32,
        }
    }
    /// Check if load balancing constraints are violated
    pub fn check_constraints(&self, stats: &LoadBalancingStats) -> bool {
        // Consider load balanced if imbalance score is below threshold
        stats.imbalance_score < 0.5 && stats.expert_utilization > 0.7
    }
 }
 /// Load balancing statistics
 #[derive(Debug, Default, Clone)]
 pub struct LoadBalancingStats {
    /// Average number of tokens per expert
    pub average_load: f32,
    /// Minimum tokens assigned to any expert
    pub min_load: f32,
    /// Maximum tokens assigned to any expert
    pub max_load: f32,
    /// Load imbalance score (0 = perfectly balanced, higher = more imbalanced)
    pub imbalance_score: f32,
    /// Fraction of experts that received at least one token
    pub expert_utilization: f32,
 }
 /// Complete Mixture of Experts layer
 #[derive(Debug)]
 pub struct MoELayer {
    /// Configuration
    config: MoEConfig,
    /// Device
    device: Device,
    /// Router for top-k gating
    router: Router,
    /// Individual expert networks
    experts: Vec<Expert>,
    /// Load balancer
    load_balancer: LoadBalancer,
 }
 impl MoELayer {
    /// Create a new MoE layer
    pub fn new(config: MoEConfig, device: &Device) -> Result<Self> {
        config.validate()?;
        // Create router
        let router = Router::new(config.clone(), device)?;
        // Create experts
        let mut experts = Vec::with_capacity(config.num_experts);
        for _ in 0..config.num_experts {
            experts.push(Expert::new(config.clone(), device)?);
        }
        // Create load balancer
        let load_balancer = LoadBalancer::new(config.clone(), device)?;
        Ok(Self {
            config,
            device: device.clone(),
            router,
            experts,
            load_balancer,
        })
    }
    /// Forward pass through the MoE layer
    pub fn forward(&self, input: &Tensor) -> Result<MoEOutput> {
        let input_shape = input.shape();
        let batch_size = input_shape[0];
        let seq_len = input_shape[1];
        // Get routing decisions
        let routing_info = self.router.route(input)?;
        // Apply capacity constraints
        let expert_capacity = self.config.calculate_expert_capacity(batch_size, seq_len);
        let routing_info = self.load_balancer.apply_capacity_constraints(&routing_info, expert_capacity)?;
        // Route tokens to experts and compute outputs
        let expert_outputs = self.compute_expert_outputs(input, &routing_info)?;
        // Combine expert outputs using routing weights
        let final_output = self.combine_expert_outputs(&expert_outputs, &routing_info)?;
        // Calculate load balancing statistics
        let load_stats = self.load_balancer.calculate_load_stats(&routing_info.expert_token_counts, batch_size * seq_len);
        Ok(MoEOutput {
            output: final_output,
            routing_info,
            load_stats,
            expert_outputs: Some(expert_outputs),
        })
    }
    /// Compute outputs from all experts
    fn compute_expert_outputs(&self, input: &Tensor, routing_info: &RoutingInfo) -> Result<Vec<Tensor>> {
        let mut expert_outputs = Vec::with_capacity(self.config.num_experts);
        // For simplicity, we compute all experts on the full input
        // In practice, we would only compute experts that have assigned tokens
        for expert in &self.experts {
            let output = expert.forward(input)?;
            expert_outputs.push(output);
        }
        Ok(expert_outputs)
    }
    /// Combine expert outputs using routing weights
    fn combine_expert_outputs(&self, expert_outputs: &[Tensor], routing_info: &RoutingInfo) -> Result<Tensor> {
        // This is a simplified implementation
        // In practice, we would:
        // 1. Scatter tokens to their assigned experts
        // 2. Compute expert outputs only on assigned tokens
        // 3. Gather and combine outputs using routing weights
        let input_shape = expert_outputs[0].shape();
        let mut combined = Tensor::zeros(input_shape, &self.device)?;
        // Weighted combination of all expert outputs
        // This is not the actual MoE implementation but serves as a placeholder
        for (i, expert_output) in expert_outputs.iter().enumerate() {
            let weight = 1.0 / self.config.num_experts as f32;
            let weighted_output = expert_output.mul_scalar(weight)?;
            combined = combined.add(&weighted_output)?;
        }
        Ok(combined)
    }
    /// Get all MoE parameters for optimization
    pub fn parameters(&self) -> Vec<&Tensor> {
        let mut params = Vec::new();
        // Router parameters
        params.extend(self.router.parameters());
        // Expert parameters
        for expert in &self.experts {
            params.extend(expert.parameters());
        }
        params
    }
    /// Get all mutable MoE parameters for optimization
    pub fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        let mut params = Vec::new();
        // Router parameters
        params.extend(self.router.parameters_mut());
        // Expert parameters
        for expert in &mut self.experts {
            params.extend(expert.parameters_mut());
        }
        params
    }
    /// Get MoE configuration
    pub fn config(&self) -> &MoEConfig {
        &self.config
    }
    /// Get number of experts
    pub fn num_experts(&self) -> usize {
        self.config.num_experts
    }
    /// Count total parameters in the MoE layer
    pub fn parameter_count(&self) -> usize {
        // Router parameters
        let router_params = self.config.hidden_dim * self.config.num_experts + 
                           if self.config.bias { self.config.num_experts } else { 0 };
        // Expert parameters
        let expert_params = self.experts[0].parameter_count() * self.config.num_experts;
        router_params + expert_params
    }
 }
 impl Layer for MoELayer {
    fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let output = self.forward(input)?;
        Ok(output.output)
    }
    fn layer_type(&self) -> &'static str {
        "MixtureOfExperts"
    }
    fn device(&self) -> &Device {
        &self.device
    }
    fn parameters(&self) -> Vec<&Tensor> {
        self.parameters()
    }
    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
        self.parameters_mut()
    }
 }
 /// Output from MoE layer forward pass
 #[derive(Debug)]
 pub struct MoEOutput {
    /// Final combined output tensor
    pub output: Tensor,
    /// Routing information
    pub routing_info: RoutingInfo,
    /// Load balancing statistics
    pub load_stats: LoadBalancingStats,
    /// Individual expert outputs (optional, for debugging)
    pub expert_outputs: Option<Vec<Tensor>>,
 }
 impl MoEOutput {
    /// Get auxiliary loss for training
    pub fn aux_loss(&self) -> Option<&Tensor> {
        self.routing_info.load_balance_loss.as_ref()
    }
    /// Check if load balancing constraints are satisfied
    pub fn is_load_balanced(&self) -> bool {
        self.load_stats.imbalance_score < 0.5
    }
 }
 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
    use super::*;
    #[test]
    fn test_load_balancer_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let load_balancer = LoadBalancer::new(config, &device);
        assert!(load_balancer.is_ok());
    }
    #[test]
    fn test_load_balancer_stats() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let load_balancer = LoadBalancer::new(config, &device).unwrap();
        // Test perfectly balanced load
        let balanced_counts = vec![16; 8]; // 8 experts, 16 tokens each
        let stats = load_balancer.calculate_load_stats(&balanced_counts, 128);
        assert_eq!(stats.average_load, 16.0);
        assert_eq!(stats.min_load, 16.0);
        assert_eq!(stats.max_load, 16.0);
        assert!(stats.imbalance_score < 1e-6);
        assert_eq!(stats.expert_utilization, 1.0);
        // Test imbalanced load
        let imbalanced_counts = vec![32, 0, 16, 8, 24, 20, 12, 16]; // Total 128
        let stats = load_balancer.calculate_load_stats(&imbalanced_counts, 128);
        assert_eq!(stats.average_load, 16.0);
        assert_eq!(stats.min_load, 0.0);
        assert_eq!(stats.max_load, 32.0);
        assert!(stats.imbalance_score > 0.0);
        assert_eq!(stats.expert_utilization, 7.0 / 8.0); // 7 experts used
    }
    #[test]
    fn test_load_balancer_constraints() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let load_balancer = LoadBalancer::new(config, &device).unwrap();
        // Well-balanced case
        let balanced_stats = LoadBalancingStats {
            average_load: 16.0,
            min_load: 14.0,
            max_load: 18.0,
            imbalance_score: 0.1,
            expert_utilization: 1.0,
        };
        assert!(load_balancer.check_constraints(&balanced_stats));
        // Poorly balanced case
        let imbalanced_stats = LoadBalancingStats {
            average_load: 16.0,
            min_load: 0.0,
            max_load: 64.0,
            imbalance_score: 1.5,
            expert_utilization: 0.5,
        };
        assert!(!load_balancer.check_constraints(&imbalanced_stats));
    }
    #[test]
    fn test_moe_layer_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let moe_layer = MoELayer::new(config.clone(), &device);
        assert!(moe_layer.is_ok());
        let moe_layer = moe_layer.unwrap();
        assert_eq!(moe_layer.num_experts(), 8);
        assert_eq!(moe_layer.config().top_k, 2);
        assert_eq!(moe_layer.config().hidden_dim, 768);
    }
    #[test]
    fn test_moe_layer_forward() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let moe_layer = MoELayer::new(config.clone(), &device).unwrap();
        // Test input: (batch_size, seq_len, hidden_dim)
        let batch_size = 4;
        let seq_len = 16;
        let input = Tensor::randn(&[batch_size, seq_len, config.hidden_dim], DType::F32, &device).unwrap();
        let output = moe_layer.forward(&input);
        assert!(output.is_ok());
        let output = output.unwrap();
        let output_shape = output.output.shape();
        // Output should have same shape as input
        assert_eq!(output_shape[0], batch_size);
        assert_eq!(output_shape[1], seq_len);
        assert_eq!(output_shape[2], config.hidden_dim);
        // Should have routing info
        assert_eq!(output.routing_info.expert_token_counts.len(), config.num_experts);
        // Should have load stats
        assert!(output.load_stats.average_load >= 0.0);
    }
    #[test]
    fn test_moe_layer_parameters() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(4, 2, 512, 2048); // Smaller config for testing
        let mut moe_layer = MoELayer::new(config.clone(), &device).unwrap();
        let params = moe_layer.parameters();
        let params_mut = moe_layer.parameters_mut();
        // Should have parameters from router + all experts
        assert!(!params.is_empty());
        assert_eq!(params.len(), params_mut.len());
        // Count should match expected
        let expected_router_params = if config.bias { 2 } else { 1 }; // weight [+ bias]
        let expected_expert_params = if config.bias { 4 } else { 2 }; // up_weight, down_weight [+ biases]
        let expected_total = expected_router_params + expected_expert_params * config.num_experts;
        assert_eq!(params.len(), expected_total);
    }
    #[test]
    fn test_moe_layer_as_layer_trait() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(8, 2, 768, 3072);
        let mut moe_layer = MoELayer::new(config.clone(), &device).unwrap();
        // Test Layer trait methods
        assert_eq!(moe_layer.layer_type(), "MixtureOfExperts");
        assert_eq!(moe_layer.device(), &device);
        // Test forward through Layer trait
        let input = Tensor::randn(&[2, 8, config.hidden_dim], DType::F32, &device).unwrap();
        let output = moe_layer.forward(&input);
        assert!(output.is_ok());
        let output = output.unwrap();
        assert_eq!(output.shape()[2], config.hidden_dim);
        // Test parameter access
        let params = moe_layer.parameters();
        let params_mut = moe_layer.parameters_mut();
        assert!(!params.is_empty());
        assert_eq!(params.len(), params_mut.len());
    }
    #[test]
    fn test_moe_output_aux_loss() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut config = MoEConfig::new(8, 2, 768, 3072);
        config.aux_loss_weight = 0.01; // Enable aux loss
        let moe_layer = MoELayer::new(config.clone(), &device).unwrap();
        let input = Tensor::randn(&[2, 8, config.hidden_dim], DType::F32, &device).unwrap();
        let output = moe_layer.forward(&input).unwrap();
        // Should have auxiliary loss when enabled
        assert!(output.aux_loss().is_some());
        // Test with disabled aux loss
        let mut config_no_aux = config.clone();
        config_no_aux.aux_loss_weight = 0.0;
        let moe_layer_no_aux = MoELayer::new(config_no_aux, &device).unwrap();
        let output_no_aux = moe_layer_no_aux.forward(&input).unwrap();
        // Should not have auxiliary loss when disabled
        assert!(output_no_aux.aux_loss().is_none());
    }
    #[test]
    fn test_parameter_count() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = MoEConfig::new(4, 2, 256, 1024);
        let moe_layer = MoELayer::new(config.clone(), &device).unwrap();
        let param_count = moe_layer.parameter_count();
        // Router: 256 * 4 (weight) + [4] (bias) = 1024 [+ 4]
        let router_params = 256 * 4 + if config.bias { 4 } else { 0 };
        // Each expert: 256 * 1024 * 2 (weights) + [256 + 1024] (biases) 
        let expert_params_each = 256 * 1024 * 2 + if config.bias { 256 + 1024 } else { 0 };
        let total_expert_params = expert_params_each * 4;
        let expected_total = router_params + total_expert_params;
        assert_eq!(param_count, expected_total);
        assert!(param_count > 0);
    }
 }
@@ -1,516 +0,0 @@
 //! Positional encoding implementations for transformer models
 //!
 //! Provides various positional encoding schemes including sinusoidal,
 //! learned, rotary, and relative positional encodings.
 use crate::{Result, TransformerError};
 use rtx_tensor::{Tensor, Shape, Device};
 use serde::{Deserialize, Serialize};
 use tracing::{debug, trace};
 use std::f64::consts::PI;
 /// Positional encoding variants for transformer models
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum PositionalEncoding {
    /// Sinusoidal positional encoding (Attention is All You Need)
    Sinusoidal {
        max_len: usize,
        d_model: usize,
        encoding: Tensor,
        device: Device,
    },
    /// Learned positional embeddings
    Learned {
        max_len: usize,
        d_model: usize,
        embeddings: Tensor,
        device: Device,
    },
    /// Rotary positional embedding (RoPE)
    Rotary {
        d_model: usize,
        max_len: usize,
        theta: f64,
        sin_cache: Tensor,
        cos_cache: Tensor,
        device: Device,
    },
    /// Relative positional encoding
    Relative {
        max_len: usize,
        d_model: usize,
        clamp_len: Option<usize>,
        embeddings: Tensor,
        device: Device,
    },
 }
 impl PositionalEncoding {
    /// Create sinusoidal positional encoding
    ///
    /// # Arguments
    /// * `max_len` - Maximum sequence length
    /// * `d_model` - Model dimension
    /// * `device` - Device to place tensors on
    pub fn sinusoidal(max_len: usize, d_model: usize, device: &Device) -> Result<Self> {
        if max_len == 0 {
            return Err(TransformerError::generic(
            "max_len",
            format!("{}", max_len),
            "must be positive",
            ));
        }
        if d_model == 0 || d_model % 2 != 0 {
            return Err(TransformerError::generic(
            "d_model",
            format!("{}", d_model),
            "must be positive and even",
            ));
        }
        debug!("Creating sinusoidal positional encoding: max_len={}, d_model={}", max_len, d_model);
        // Create position indices [0, 1, 2, ..., max_len-1]
        let positions: Vec<f64> = (0..max_len).map(|i| i as f64).collect();
        let position_tensor = Tensor::from_slice(&positions, &Shape::new(vec![max_len, 1]), device)?;
        // Create dimension indices [0, 2, 4, ..., d_model-2]
        let dim_indices: Vec<f64> = (0..d_model/2).map(|i| i as f64).collect();
        let dim_tensor = Tensor::from_slice(&dim_indices, &Shape::new(vec![1, d_model/2]), device)?;
        // Compute div_term = 1 / (10000 ^ (2i / d_model))
        let div_term = (&dim_tensor * (2.0 / d_model as f64))?.exp()? * (-2.0 * PI.ln());
        let div_term = div_term.exp()?;
        // Compute angles = position / div_term
        let angles = &position_tensor / &div_term;
        // Create encoding tensor
        let mut encoding_data = Vec::with_capacity(max_len * d_model);
        // Fill with sin and cos alternating
        for pos in 0..max_len {
            for i in 0..d_model/2 {
                let angle = angles.get(&[pos, i])?.to_scalar::<f64>()?;
                encoding_data.push(angle.sin() as f32);
                encoding_data.push(angle.cos() as f32);
            }
        }
        let encoding = Tensor::from_slice(&encoding_data, &Shape::new(vec![max_len, d_model]), device)?;
        Ok(Self::Sinusoidal {
            max_len,
            d_model,
            encoding,
            device: device.clone(),
        })
    }
    /// Create learned positional embeddings
    ///
    /// # Arguments
    /// * `max_len` - Maximum sequence length
    /// * `d_model` - Model dimension
    /// * `device` - Device to place tensors on
    pub fn learned(max_len: usize, d_model: usize, device: &Device) -> Result<Self> {
        if max_len == 0 {
            return Err(TransformerError::generic(
            "max_len",
            format!("{}", max_len),
            "must be positive",
            ));
        }
        if d_model == 0 {
            return Err(TransformerError::generic(
            "d_model",
            format!("{}", d_model),
            "must be positive",
            ));
        }
        debug!("Creating learned positional encoding: max_len={}, d_model={}", max_len, d_model);
        // Initialize with small random values
        let embeddings = Tensor::randn(&Shape::new(vec![max_len, d_model]), device)? * 0.1;
        Ok(Self::Learned {
            max_len,
            d_model,
            embeddings,
            device: device.clone(),
        })
    }
    /// Create rotary positional embedding (RoPE)
    ///
    /// # Arguments
    /// * `d_model` - Model dimension (must be even)
    /// * `max_len` - Maximum sequence length
    /// * `theta` - Base for frequency computation (typically 10000.0)
    /// * `device` - Device to place tensors on
    pub fn rotary(d_model: usize, max_len: usize, theta: f64, device: &Device) -> Result<Self> {
        if d_model == 0 || d_model % 2 != 0 {
            return Err(TransformerError::generic(
            "d_model",
            format!("{}", d_model),
            "must be positive and even",
            ));
        }
        if max_len == 0 {
            return Err(TransformerError::generic(
            "max_len",
            format!("{}", max_len),
            "must be positive",
            ));
        }
        if theta <= 0.0 {
            return Err(TransformerError::generic(
            "theta",
            format!("{}", theta),
            "must be positive",
            ));
        }
        debug!("Creating rotary positional encoding: d_model={}, max_len={}, theta={}", d_model, max_len, theta);
        // Compute frequencies
        let half_dim = d_model / 2;
        let freqs: Vec<f64> = (0..half_dim)
        .map(|i| 1.0 / theta.powf(2.0 * i as f64 / d_model as f64))
        .collect();
        // Create position indices
        let positions: Vec<f64> = (0..max_len).map(|i| i as f64).collect();
        // Compute angles for each position and frequency
        let mut sin_data = Vec::with_capacity(max_len * half_dim);
        let mut cos_data = Vec::with_capacity(max_len * half_dim);
        for &pos in &positions {
            for &freq in &freqs {
                let angle = pos * freq;
                sin_data.push(angle.sin() as f32);
                cos_data.push(angle.cos() as f32);
            }
        }
        let sin_cache = Tensor::from_slice(&sin_data, &Shape::new(vec![max_len, half_dim]), device)?;
        let cos_cache = Tensor::from_slice(&cos_data, &Shape::new(vec![max_len, half_dim]), device)?;
        Ok(Self::Rotary {
            d_model,
            max_len,
            theta,
            sin_cache,
            cos_cache,
            device: device.clone(),
        })
    }
    /// Create relative positional encoding
    ///
    /// # Arguments
    /// * `max_len` - Maximum sequence length
    /// * `d_model` - Model dimension
    /// * `clamp_len` - Optional maximum relative distance
    /// * `device` - Device to place tensors on
    pub fn relative(max_len: usize, d_model: usize, clamp_len: Option<usize>, device: &Device) -> Result<Self> {
        if max_len == 0 {
            return Err(TransformerError::generic(
            "max_len",
            format!("{}", max_len),
            "must be positive",
            ));
        }
        if d_model == 0 {
            return Err(TransformerError::generic(
            "d_model",
            format!("{}", d_model),
            "must be positive",
            ));
        }
        debug!("Creating relative positional encoding: max_len={}, d_model={}, clamp_len={:?}", max_len, d_model, clamp_len);
        // Relative positions range from -(max_len-1) to +(max_len-1)
        let relative_range = 2 * max_len - 1;
        let embeddings = Tensor::randn(&Shape::new(vec![relative_range, d_model]), device)? * 0.1;
        Ok(Self::Relative {
            max_len,
            d_model,
            clamp_len,
            embeddings,
            device: device.clone(),
        })
    }
    /// Get the maximum sequence length
    pub fn max_len(&self) -> usize {
        match self {
            Self::Sinusoidal { max_len, .. } => *max_len,
            Self::Learned { max_len, .. } => *max_len,
            Self::Rotary { max_len, .. } => *max_len,
            Self::Relative { max_len, .. } => *max_len,
        }
    }
    /// Get the model dimension
    pub fn d_model(&self) -> usize {
        match self {
            Self::Sinusoidal { d_model, .. } => *d_model,
            Self::Learned { d_model, .. } => *d_model,
            Self::Rotary { d_model, .. } => *d_model,
            Self::Relative { d_model, .. } => *d_model,
        }
    }
    /// Apply positional encoding to input
    pub fn apply(&self, input: &Tensor, start_pos: usize) -> Result<Tensor> {
        match self {
            Self::Sinusoidal { encoding, .. } => {
                self.apply_additive(input, encoding, start_pos)
            }
            Self::Learned { embeddings, .. } => {
                self.apply_additive(input, embeddings, start_pos)
            }
            Self::Rotary { sin_cache, cos_cache, .. } => {
                self.apply_rotary(input, sin_cache, cos_cache, start_pos)
            }
            Self::Relative { .. } => {
                // Relative encoding requires special handling in attention computation
                // For now, return input unchanged
                Ok(input.clone())
            }
        }
    }
    /// Apply additive positional encoding
    fn apply_additive(&self, input: &Tensor, encoding: &Tensor, start_pos: usize) -> Result<Tensor> {
        let input_shape = input.shape();
        let input_dims = input_shape.dims();
        if input_dims.len() < 2 {
            return Err(TransformerError::dimension(
            "Input must have at least 2 dimensions for positional encoding".to_string()
            ));
        }
        let seq_len = input_dims[input_dims.len() - 2];
        let model_dim = input_dims[input_dims.len() - 1];
        if model_dim != self.d_model() {
            return Err(TransformerError::shape_mismatch(
            vec![model_dim],
            vec![self.d_model()],
            ));
        }
        if start_pos + seq_len > self.max_len() {
            return Err(TransformerError::generic(
            "sequence_position",
            format!("start_pos={}, seq_len={}", start_pos, seq_len),
            format!("exceeds max_len={}", self.max_len()),
            ));
        }
        // Extract relevant portion of encoding
        let pos_encoding = encoding.slice(&[start_pos..start_pos + seq_len, 0..model_dim])?;
        // Add positional encoding to input
        let output = input + &pos_encoding;
        trace!("Applied additive positional encoding: start_pos={}, seq_len={}", start_pos, seq_len);
        Ok(output)
    }
    /// Apply rotary positional encoding
    fn apply_rotary(&self, input: &Tensor, sin_cache: &Tensor, cos_cache: &Tensor, start_pos: usize) -> Result<Tensor> {
        let input_shape = input.shape();
        let input_dims = input_shape.dims();
        if input_dims.len() < 2 {
            return Err(TransformerError::dimension(
            "Input must have at least 2 dimensions for rotary encoding".to_string()
            ));
        }
        let seq_len = input_dims[input_dims.len() - 2];
        let model_dim = input_dims[input_dims.len() - 1];
        if model_dim != self.d_model() {
            return Err(TransformerError::shape_mismatch(
            vec![model_dim],
            vec![self.d_model()],
            ));
        }
        if start_pos + seq_len > self.max_len() {
            return Err(TransformerError::generic(
            "sequence_position",
            format!("start_pos={}, seq_len={}", start_pos, seq_len),
            format!("exceeds max_len={}", self.max_len()),
            ));
        }
        // Split input into x1 and x2 (first and second half)
        let half_dim = model_dim / 2;
        let x1 = input.slice(&[.., 0..half_dim])?;
        let x2 = input.slice(&[.., half_dim..model_dim])?;
        // Get sin and cos for current positions
        let sin_pos = sin_cache.slice(&[start_pos..start_pos + seq_len, ..])?;
        let cos_pos = cos_cache.slice(&[start_pos..start_pos + seq_len, ..])?;
        // Apply rotary transformation: [x1*cos - x2*sin, x1*sin + x2*cos]
        let rotated_x1 = &x1 * &cos_pos - &x2 * &sin_pos;
        let rotated_x2 = &x1 * &sin_pos + &x2 * &cos_pos;
        // Concatenate rotated halves
        let output = Tensor::concat(&[rotated_x1, rotated_x2], input_dims.len() - 1)?;
        trace!("Applied rotary positional encoding: start_pos={}, seq_len={}", start_pos, seq_len);
        Ok(output)
    }
    /// Get relative position bias for attention computation
    pub fn get_relative_bias(&self, query_len: usize, key_len: usize) -> Result<Option<Tensor>> {
        match self {
            Self::Relative { embeddings, max_len, clamp_len, .. } => {
                let mut bias_data = Vec::with_capacity(query_len * key_len);
                for i in 0..query_len {
                    for j in 0..key_len {
                        let relative_pos = (i as i32 - j as i32) + (*max_len as i32 - 1);
                        let clamped_pos = if let Some(clamp) = clamp_len {
                            relative_pos.max(0).min(*clamp as i32 - 1)
                            } else {
                                relative_pos.max(0).min((2 * max_len - 2) as i32)
                            };
                            bias_data.push(clamped_pos as f32);
                        }
                    }
                    let indices = Tensor::from_slice(&bias_data, &Shape::new(vec![query_len, key_len]), self.device())?;
                    let bias = embeddings.index_select(0, &indices)?;
                    Ok(Some(bias))
                }
                _ => Ok(None),
            }
        }
    }
    impl Layer for PositionalEncoding {
        fn forward(&self, input: &Tensor) -> Result<Tensor> {
            self.apply(input, 0)
        }
        fn layer_type(&self) -> &'static str {
            match self {
                Self::Sinusoidal { .. } => "SinusoidalPositionalEncoding",
                Self::Learned { .. } => "LearnedPositionalEncoding",
                Self::Rotary { .. } => "RotaryPositionalEncoding",
                Self::Relative { .. } => "RelativePositionalEncoding",
            }
        }
        fn device(&self) -> &Device {
            match self {
                Self::Sinusoidal { device, .. } => device,
                Self::Learned { device, .. } => device,
                Self::Rotary { device, .. } => device,
                Self::Relative { device, .. } => device,
            }
        }
        fn parameters(&self) -> Vec<&Tensor> {
            match self {
                Self::Sinusoidal { .. } => vec![], // Fixed encoding
                Self::Learned { embeddings, .. } => vec![embeddings],
                Self::Rotary { .. } => vec![], // Fixed encoding
                Self::Relative { embeddings, .. } => vec![embeddings],
            }
        }
        fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
            match self {
                Self::Sinusoidal { .. } => vec![], // Fixed encoding
                Self::Learned { embeddings, .. } => vec![embeddings],
                Self::Rotary { .. } => vec![], // Fixed encoding
                Self::Relative { embeddings, .. } => vec![embeddings],
            }
        }
    }
    #[cfg(test)]
    mod tests {
        use super::*;
        use rtx_tensor::Device;
        #[tokio::test]
        async fn test_sinusoidal_encoding() -> Result<()> {
            let device = Device::cuda(0)?;
            let pos_enc = PositionalEncoding::sinusoidal(512, 128, &device)?;
            assert_eq!(pos_enc.max_len(), 512);
            assert_eq!(pos_enc.d_model(), 128);
            Ok(())
        }
        #[tokio::test]
        async fn test_learned_encoding() -> Result<()> {
            let device = Device::cuda(0)?;
            let pos_enc = PositionalEncoding::learned(512, 128, &device)?;
            assert_eq!(pos_enc.max_len(), 512);
            assert_eq!(pos_enc.d_model(), 128);
            assert_eq!(pos_enc.parameters().len(), 1);
            Ok(())
        }
        #[tokio::test]
        async fn test_rotary_encoding() -> Result<()> {
            let device = Device::cuda(0)?;
            let pos_enc = PositionalEncoding::rotary(128, 512, 10000.0, &device)?;
            assert_eq!(pos_enc.max_len(), 512);
            assert_eq!(pos_enc.d_model(), 128);
            Ok(())
        }
        #[tokio::test]
        async fn test_relative_encoding() -> Result<()> {
            let device = Device::cuda(0)?;
            let pos_enc = PositionalEncoding::relative(512, 128, Some(64), &device)?;
            assert_eq!(pos_enc.max_len(), 512);
            assert_eq!(pos_enc.d_model(), 128);
            assert_eq!(pos_enc.parameters().len(), 1);
            Ok(())
        }
        #[test]
        fn test_invalid_parameters() {
            let device = Device::cuda(0).unwrap();
            // Invalid d_model for sinusoidal (must be even)
            assert!(PositionalEncoding::sinusoidal(512, 127, &device).is_err());
            // Invalid max_len
            assert!(PositionalEncoding::learned(0, 128, &device).is_err());
            // Invalid theta for rotary
            assert!(PositionalEncoding::rotary(128, 512, 0.0, &device).is_err());
        }
    }
@@ -1,871 +0,0 @@
 //! Sliding Window Attention Implementation
 //!
 //! This module implements Sliding Window Attention, which limits attention to a 
 //! local window of tokens rather than the full sequence. This approach enables
 //! efficient processing of long sequences by reducing computational complexity
 //! from O(n²) to O(n*w) where w is the window size.
 //!
 //! ## Features
 //!
 //! - **Configurable window size**: Support for different window sizes (e.g., 256 tokens)
 //! - **Efficient computation**: O(n*w) complexity for long sequences instead of O(n²)
 //! - **Boundary handling**: Proper handling at sequence start/end
 //! - **Causal and bidirectional modes**: Support for both autoregressive and encoder models
 //! - **MQA/GQA integration**: Full compatibility with Multi-Query and Grouped-Query Attention
 //! - **Memory efficiency**: Significant memory savings for long sequences
 //!
 //! ## Usage
 //!
 //! ```rust
 //! use rtx_transformers::layers::{SlidingWindowAttention, SlidingWindowConfig};
 //! use rtx_tensor::Device;
 //!
 //! // Create configuration
 //! let config = SlidingWindowConfig::new(12, 768, 256)?; // 12 heads, 768 dim, 256 window
 //! let device = Device::cuda(0).unwrap_or(Device::default());
 //!
 //! // Create attention layer
 //! let mut attention = SlidingWindowAttention::new(config, &device)?;
 //! attention.initialize_parameters()?;
 //!
 //! // Use in forward pass
 //! let output = attention.forward(&input_tensor, None, None)?;
 //! ```
 //!
 //! ## Integration with MQA/GQA
 //!
 //! ```rust
 //! // Multi-Query Attention (1 KV head)
 //! let mqa_config = SlidingWindowConfig::new(12, 768, 256)?
 //!     .with_kv_heads(1)?;
 //!
 //! // Grouped-Query Attention (4 KV heads)  
 //! let gqa_config = SlidingWindowConfig::new(12, 768, 256)?
 //!     .with_kv_heads(4)?;
 //! ```
 use crate::{Result, TransformerError};
 use crate::architectures::TransformerConfig;
 use rtx_tensor::{Tensor, Device, DType};
 use rtx_autograd::{TensorAutograd, backward, NodeId};
 use serde::{Deserialize, Serialize};
 /// Sliding Window Attention configuration
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SlidingWindowConfig {
    /// Number of attention heads
    pub num_heads: usize,
    /// Model dimension
    pub d_model: usize,
    /// Head dimension (d_model / num_heads)
    pub head_dim: usize,
    /// Window size (number of tokens to attend to on each side)
    pub window_size: usize,
    /// Whether to use causal masking (only attend to previous tokens)
    pub causal: bool,
    /// Whether to apply scaling factor (1/sqrt(head_dim))
    pub scale: bool,
    /// Dropout probability for attention weights
    pub dropout: f64,
    /// Whether to use bias in linear projections
    pub bias: bool,
    /// Number of key-value heads (for integration with MQA/GQA)
    pub num_kv_heads: Option<usize>,
 }
 impl SlidingWindowConfig {
    /// Create sliding window config from TransformerConfig
    pub fn from_transformer_config(config: &TransformerConfig, window_size: usize) -> Result<Self> {
        if window_size == 0 {
            return Err(TransformerError::config(
                "Window size must be greater than 0".to_string()
            ));
        }
        let head_dim = config.d_model / config.num_heads;
        if config.d_model % config.num_heads != 0 {
            return Err(TransformerError::config(
                format!("d_model ({}) must be divisible by num_heads ({})", 
                       config.d_model, config.num_heads)
            ));
        }
        let num_kv_heads = if config.get_num_key_value_heads() != config.num_heads {
            Some(config.get_num_key_value_heads())
        } else {
            None
        };
        Ok(Self {
            num_heads: config.num_heads,
            d_model: config.d_model,
            head_dim,
            window_size,
            causal: true, // Default to causal for autoregressive models
            scale: true,
            dropout: config.dropout,
            bias: config.bias,
            num_kv_heads,
        })
    }
    /// Create config with specific parameters
    pub fn new(
        num_heads: usize,
        d_model: usize,
        window_size: usize,
    ) -> Result<Self> {
        if window_size == 0 {
            return Err(TransformerError::config(
                "Window size must be greater than 0".to_string()
            ));
        }
        if d_model % num_heads != 0 {
            return Err(TransformerError::config(
                format!("d_model ({}) must be divisible by num_heads ({})", 
                       d_model, num_heads)
            ));
        }
        Ok(Self {
            num_heads,
            d_model,
            head_dim: d_model / num_heads,
            window_size,
            causal: true,
            scale: true,
            dropout: 0.0,
            bias: true,
            num_kv_heads: None,
        })
    }
    /// Enable/disable causal masking
    pub fn with_causal(mut self, causal: bool) -> Self {
        self.causal = causal;
        self
    }
    /// Set number of key-value heads for MQA/GQA integration
    pub fn with_kv_heads(mut self, num_kv_heads: usize) -> Result<Self> {
        if num_kv_heads > self.num_heads {
            return Err(TransformerError::config(
                format!("num_kv_heads ({}) cannot exceed num_heads ({})", 
                       num_kv_heads, self.num_heads)
            ));
        }
        if self.num_heads % num_kv_heads != 0 {
            return Err(TransformerError::config(
                format!("num_heads ({}) must be divisible by num_kv_heads ({})", 
                       self.num_heads, num_kv_heads)
            ));
        }
        self.num_kv_heads = Some(num_kv_heads);
        Ok(self)
    }
    /// Get effective number of key-value heads
    pub fn get_num_kv_heads(&self) -> usize {
        self.num_kv_heads.unwrap_or(self.num_heads)
    }
    /// Check if this is Multi-Query Attention (single KV head)
    pub fn is_mqa(&self) -> bool {
        self.num_kv_heads == Some(1)
    }
    /// Check if this is Grouped-Query Attention (multiple but fewer KV heads)
    pub fn is_gqa(&self) -> bool {
        match self.num_kv_heads {
            Some(kv_heads) => kv_heads > 1 && kv_heads < self.num_heads,
            None => false,
        }
    }
 }
 /// Sliding Window Attention layer
 #[derive(Debug)]
 pub struct SlidingWindowAttention {
    config: SlidingWindowConfig,
    device: Device,
    // Placeholder for actual tensor parameters
    q_proj_initialized: bool,
    k_proj_initialized: bool,
    v_proj_initialized: bool,
    o_proj_initialized: bool,
 }
 impl SlidingWindowAttention {
    /// Create new Sliding Window Attention layer
    pub fn new(config: SlidingWindowConfig, device: &Device) -> Result<Self> {
        Ok(Self {
            config,
            device: device.clone(),
            q_proj_initialized: false,
            k_proj_initialized: false,
            v_proj_initialized: false,
            o_proj_initialized: false,
        })
    }
    /// Initialize layer parameters
    pub fn initialize_parameters(&mut self) -> Result<()> {
        // In a real implementation, this would initialize:
        // - q_proj: Linear(d_model, num_heads * head_dim)
        // - k_proj: Linear(d_model, num_kv_heads * head_dim)
        // - v_proj: Linear(d_model, num_kv_heads * head_dim)
        // - o_proj: Linear(num_heads * head_dim, d_model)
        self.q_proj_initialized = true;
        self.k_proj_initialized = true;
        self.v_proj_initialized = true;
        self.o_proj_initialized = true;
        Ok(())
    }
    /// Forward pass with sliding window attention
    pub fn forward(
        &self,
        hidden_states: &Tensor,
        attention_mask: Option<&Tensor>,
        position_ids: Option<&Tensor>,
    ) -> Result<Tensor> {
        if !self.q_proj_initialized {
            return Err(TransformerError::architecture(
                "Sliding window attention layer parameters not initialized. Call initialize_parameters() first.".to_string()
            ));
        }
        let batch_size = hidden_states.shape().dims()[0];
        let seq_len = hidden_states.shape().dims()[1];
        let d_model = hidden_states.shape().dims()[2];
        if d_model != self.config.d_model {
            return Err(TransformerError::shape_mismatch(
                format!("Expected d_model {}, got {}", self.config.d_model, d_model)
            ));
        }
        // Implement sliding window attention computation
        // This is a simplified but functional implementation that demonstrates
        // the key concepts of sliding window attention
        let scale_factor = if self.config.scale {
            1.0 / (self.config.head_dim as f32).sqrt()
        } else {
            1.0
        };
        // Step 1: Simulate query, key, value projections
        let num_kv_heads = self.config.get_num_kv_heads();
        // Query projection: [batch, seq_len, num_heads * head_dim]
        let query_scale = (self.config.num_heads as f32 / d_model as f32).sqrt() * scale_factor;
        let queries = hidden_states.mul_scalar(query_scale)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
        // Key/Value projection: [batch, seq_len, num_kv_heads * head_dim]  
        let kv_scale = (num_kv_heads as f32 / d_model as f32).sqrt() * scale_factor;
        let keys = hidden_states.mul_scalar(kv_scale * 0.9)  // Slightly different from queries
            .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
        let values = hidden_states.mul_scalar(kv_scale * 1.1)  // Different scaling for values
            .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
        // Step 2: Apply sliding window attention pattern
        let effective_window = self.effective_attention_span(seq_len);
        let window_factor = (effective_window as f32) / (seq_len as f32);
        // Step 3: Simulate attention computation with window constraints
        // In a real implementation, this would involve:
        // - Reshaping Q, K, V to separate heads
        // - Computing attention scores only within the window
        // - Applying causal/bidirectional masking as needed
        // - Applying attention mask if provided
        let attention_output = if self.config.causal {
            // Causal sliding window: each token can only attend to previous tokens within window
            let causal_factor = 0.8 + window_factor * 0.2;
            queries.mul_scalar(causal_factor)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
                .add(&values.mul_scalar(0.3)
                    .map_err(|e| TransformerError::tensor_op(e.to_string()))?)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
        } else {
            // Bidirectional sliding window: can attend to tokens on both sides
            let bidirectional_factor = 0.7 + window_factor * 0.3;
            queries.mul_scalar(bidirectional_factor)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
                .add(&values.mul_scalar(0.4)
                    .map_err(|e| TransformerError::tensor_op(e.to_string()))?)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
        };
        // Step 4: Apply output projection
        // In a real implementation: output = attention_output @ W_o
        let output_factor = if num_kv_heads < self.config.num_heads { 
            // Account for reduced KV heads in MQA/GQA
            1.1 - (num_kv_heads as f32 / self.config.num_heads as f32) * 0.1
        } else { 
            1.0 
        };
        let final_output = attention_output.mul_scalar(output_factor)
            .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
        // Step 5: Apply attention mask if provided
        let masked_output = if let Some(mask) = attention_mask {
            // Simple mask application - in real implementation would be more sophisticated
            let mask_expanded = mask.unsqueeze(2)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
                .broadcast_to(&final_output.shape().dims())
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
            final_output.mul(&mask_expanded)
                .map_err(|e| TransformerError::tensor_op(e.to_string()))?
        } else {
            final_output
        };
        Ok(masked_output)
    }
    /// Get configuration
    pub fn config(&self) -> &SlidingWindowConfig {
        &self.config
    }
    /// Compute effective attention span for given sequence length
    pub fn effective_attention_span(&self, seq_len: usize) -> usize {
        if self.config.causal {
            // Causal attention: can only attend to previous tokens within window
            (self.config.window_size + 1).min(seq_len)
        } else {
            // Bidirectional: can attend to tokens on both sides within window
            (2 * self.config.window_size + 1).min(seq_len)
        }
    }
    /// Estimate memory savings compared to full attention
    pub fn memory_savings_ratio(&self, seq_len: usize) -> f64 {
        let effective_span = self.effective_attention_span(seq_len) as f64;
        let full_span = seq_len as f64;
        // Memory scales with attention matrix size
        let sliding_memory = effective_span;
        let full_memory = full_span;
        1.0 - (sliding_memory / full_memory)
    }
 }
 #[cfg(all(test, feature = "disabled_tests"))]
 mod tests {
    use super::*;
    use crate::architectures::TransformerConfig;
    #[test]
    fn test_sliding_window_config_creation_should_fail_initially() {
        // This test should FAIL initially - testing basic config creation
        let result = SlidingWindowConfig::new(12, 768, 256);
        // This will fail because SlidingWindowConfig::new doesn't exist yet
        assert!(result.is_ok());
        let config = result.unwrap();
        assert_eq!(config.num_heads, 12);
        assert_eq!(config.d_model, 768);
        assert_eq!(config.window_size, 256);
        assert_eq!(config.head_dim, 64); // 768 / 12
        assert!(config.causal); // Default should be causal
    }
    #[test]
    fn test_sliding_window_config_validation_should_fail_initially() {
        // This test should FAIL initially - testing validation
        let result = SlidingWindowConfig::new(12, 768, 0);
        // Should fail with zero window size
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Window size must be greater than 0"));
    }
    #[test]
    fn test_sliding_window_config_from_transformer_config_should_fail_initially() {
        // This test should FAIL initially - testing creation from transformer config
        let transformer_config = TransformerConfig::new(50257, 768, 12, 12, 3072, 1024);
        let result = SlidingWindowConfig::from_transformer_config(&transformer_config, 256);
        assert!(result.is_ok());
        let config = result.unwrap();
        assert_eq!(config.window_size, 256);
        assert_eq!(config.num_heads, 12);
    }
    #[test]
    fn test_sliding_window_config_causal_settings_should_fail_initially() {
        // This test should FAIL initially - testing causal configuration
        let config = SlidingWindowConfig::new(8, 512, 128).unwrap()
            .with_causal(false);
        assert!(!config.causal); // Should be bidirectional
        assert_eq!(config.window_size, 128);
    }
    #[test]
    fn test_sliding_window_config_mqa_integration_should_fail_initially() {
        // This test should FAIL initially - testing MQA integration
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap()
            .with_kv_heads(1).unwrap(); // Multi-Query Attention
        assert!(config.is_mqa());
        assert!(!config.is_gqa());
        assert_eq!(config.get_num_kv_heads(), 1);
    }
    #[test]
    fn test_sliding_window_config_gqa_integration_should_fail_initially() {
        // This test should FAIL initially - testing GQA integration
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap()
            .with_kv_heads(4).unwrap(); // Grouped-Query Attention
        assert!(!config.is_mqa());
        assert!(config.is_gqa());
        assert_eq!(config.get_num_kv_heads(), 4);
    }
    #[test]
    fn test_sliding_window_attention_creation_should_fail_initially() {
        // This test should FAIL initially - testing attention layer creation
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let attention = SlidingWindowAttention::new(config, &device);
        assert!(attention.is_ok());
    }
    #[test]
    fn test_sliding_window_attention_forward_pass_should_fail_initially() {
        // This test should FAIL initially - testing forward pass
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        // Create test input
        let batch_size = 2;
        let seq_len = 512; // Longer than window size
        let d_model = 768;
        let hidden_states = Tensor::zeros_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, None, None);
        assert!(result.is_ok());
        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_model]);
    }
    #[test]
    fn test_sliding_window_effective_attention_span_should_fail_initially() {
        // This test should FAIL initially - testing attention span calculation
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let attention = SlidingWindowAttention::new(config, &device).unwrap();
        // Test causal attention span
        let causal_span = attention.effective_attention_span(512);
        assert_eq!(causal_span, 257); // window_size + 1 for causal
        // Test with sequence shorter than window
        let short_span = attention.effective_attention_span(128);
        assert_eq!(short_span, 128); // Limited by sequence length
    }
    #[test]
    fn test_sliding_window_bidirectional_attention_span_should_fail_initially() {
        // This test should FAIL initially - testing bidirectional attention span
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap()
            .with_causal(false);
        let device = Device::cuda(0).unwrap_or(Device::default());
        let attention = SlidingWindowAttention::new(config, &device).unwrap();
        // Test bidirectional attention span
        let bidirectional_span = attention.effective_attention_span(512);
        assert_eq!(bidirectional_span, 513); // 2 * window_size + 1
        // Test with sequence shorter than window
        let short_span = attention.effective_attention_span(128);
        assert_eq!(short_span, 128); // Limited by sequence length
    }
    #[test]
    fn test_sliding_window_memory_savings_should_fail_initially() {
        // This test should FAIL initially - testing memory savings calculation
        let config = SlidingWindowConfig::new(12, 768, 256).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let attention = SlidingWindowAttention::new(config, &device).unwrap();
        // Test memory savings for long sequence
        let savings = attention.memory_savings_ratio(2048);
        assert!(savings > 0.8); // Should save significant memory
        assert!(savings < 1.0); // But not 100%
        // Test with sequence equal to window size
        let no_savings = attention.memory_savings_ratio(257); // window_size + 1
        assert!(no_savings < 0.1); // Minimal savings when sequence fits in window
    }
    #[test]
    fn test_sliding_window_long_sequence_forward_should_fail_initially() {
        // This test should FAIL initially - testing forward pass with long sequences
        let config = SlidingWindowConfig::new(8, 512, 64).unwrap(); // Small window for testing
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        // Create a long sequence that exceeds the window size
        let batch_size = 1;
        let seq_len = 256; // Much longer than window_size=64
        let d_model = 512;
        let hidden_states = Tensor::ones_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, None, None);
        assert!(result.is_ok());
        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_model]);
        // Output should be different from input (showing attention computation)
        let input_sum = hidden_states.sum(None).unwrap().to_scalar::<f32>().unwrap();
        let output_sum = output.sum(None).unwrap().to_scalar::<f32>().unwrap();
        assert_ne!(input_sum, output_sum, "Output should differ from input due to attention");
    }
    #[test]
    fn test_sliding_window_gradient_computation_should_fail_initially() {
        // This test should FAIL initially - testing gradient computation through sliding window
        let config = SlidingWindowConfig::new(4, 256, 32).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        let batch_size = 1;
        let seq_len = 64;
        let d_model = 256;
        let input_data = vec![0.5f32; batch_size * seq_len * d_model];
        let hidden_states = Tensor::from_vec(
            input_data,
            &[batch_size, seq_len, d_model],
            &device
        ).unwrap();
        let hidden_states_grad = hidden_states.set_requires_grad(true);
        // Forward pass
        let output = attention.forward(&hidden_states_grad, None, None).unwrap();
        // Create loss
        let loss = output.sum(None).unwrap();
        // Backward pass
        let grad_result = backward(&loss, &[&hidden_states_grad]);
        assert!(grad_result.is_ok());
        let gradients = grad_result.unwrap();
        assert_eq!(gradients.len(), 1);
        let input_grad = &gradients[0];
        assert_eq!(input_grad.shape().dims(), hidden_states_grad.shape().dims());
        // Gradients should be meaningful (not all zeros)
        let grad_norm = input_grad.sum(None).unwrap().to_scalar::<f32>().unwrap();
        assert!(grad_norm.abs() > 1e-6, "Gradients should be non-zero for proper backprop");
    }
    #[test] 
    fn test_sliding_window_causal_vs_bidirectional_should_fail_initially() {
        // This test should FAIL initially - testing causal vs bidirectional modes
        let device = Device::cuda(0).unwrap_or(Device::default());
        let window_size = 32;
        // Create causal config
        let causal_config = SlidingWindowConfig::new(8, 512, window_size).unwrap()
            .with_causal(true);
        let causal_attention = SlidingWindowAttention::new(causal_config, &device).unwrap();
        // Create bidirectional config  
        let bidirectional_config = SlidingWindowConfig::new(8, 512, window_size).unwrap()
            .with_causal(false);
        let bidirectional_attention = SlidingWindowAttention::new(bidirectional_config, &device).unwrap();
        let seq_len = 128;
        // Test attention spans
        let causal_span = causal_attention.effective_attention_span(seq_len);
        let bidirectional_span = bidirectional_attention.effective_attention_span(seq_len);
        // Bidirectional should have larger effective span
        assert!(bidirectional_span > causal_span);
        assert_eq!(causal_span, window_size + 1); // window_size + current token
        assert_eq!(bidirectional_span, 2 * window_size + 1); // both directions + current token
    }
    #[test]
    fn test_sliding_window_with_attention_mask_should_fail_initially() {
        // This test should FAIL initially - testing interaction with attention masks
        let config = SlidingWindowConfig::new(4, 256, 16).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        let batch_size = 2;
        let seq_len = 32;
        let d_model = 256;
        let hidden_states = Tensor::ones_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        // Create attention mask (e.g., padding mask)
        let mask_data = vec![1.0f32; batch_size * seq_len];
        let attention_mask = Tensor::from_vec(
            mask_data,
            &[batch_size, seq_len],
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, Some(&attention_mask), None);
        assert!(result.is_ok());
        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_model]);
    }
    #[test]
    fn test_sliding_window_boundary_tokens_should_fail_initially() {
        // This test should FAIL initially - testing behavior at sequence boundaries
        let config = SlidingWindowConfig::new(4, 128, 8).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        // Test with sequence exactly equal to window size
        let seq_len = 8; // Same as window_size
        let batch_size = 1;
        let d_model = 128;
        let hidden_states = Tensor::ones_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, None, None);
        assert!(result.is_ok());
        // Test with sequence smaller than window
        let small_seq_len = 4; // Smaller than window_size=8
        let small_hidden_states = Tensor::ones_typed(
            &[batch_size, small_seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let small_result = attention.forward(&small_hidden_states, None, None);
        assert!(small_result.is_ok());
        let small_output = small_result.unwrap();
        assert_eq!(small_output.shape().dims(), &[batch_size, small_seq_len, d_model]);
    }
    #[test]
    fn test_sliding_window_with_mqa_integration_should_fail_initially() {
        // This test should FAIL initially - testing MQA integration
        let base_config = SlidingWindowConfig::new(12, 768, 128).unwrap();
        let mqa_config = base_config.with_kv_heads(1).unwrap(); // Multi-Query Attention
        assert!(mqa_config.is_mqa());
        assert!(!mqa_config.is_gqa());
        assert_eq!(mqa_config.get_num_kv_heads(), 1);
        assert_eq!(mqa_config.num_heads, 12);
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(mqa_config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        let batch_size = 1;
        let seq_len = 64;
        let d_model = 768;
        let hidden_states = Tensor::ones_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, None, None);
        assert!(result.is_ok());
        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_model]);
    }
    #[test]
    fn test_sliding_window_with_gqa_integration_should_fail_initially() {
        // This test should FAIL initially - testing GQA integration
        let base_config = SlidingWindowConfig::new(12, 768, 128).unwrap();
        let gqa_config = base_config.with_kv_heads(4).unwrap(); // Grouped-Query Attention
        assert!(!gqa_config.is_mqa());
        assert!(gqa_config.is_gqa());
        assert_eq!(gqa_config.get_num_kv_heads(), 4);
        assert_eq!(gqa_config.num_heads, 12);
        let device = Device::cuda(0).unwrap_or(Device::default());
        let mut attention = SlidingWindowAttention::new(gqa_config, &device).unwrap();
        attention.initialize_parameters().unwrap();
        let batch_size = 2;
        let seq_len = 96;
        let d_model = 768;
        let hidden_states = Tensor::ones_typed(
            &[batch_size, seq_len, d_model],
            DType::F32,
            &device
        ).unwrap();
        let result = attention.forward(&hidden_states, None, None);
        assert!(result.is_ok());
        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_model]);
    }
    #[test]
    fn test_sliding_window_kv_heads_validation_should_fail_initially() {
        // This test should FAIL initially - testing KV heads validation
        let base_config = SlidingWindowConfig::new(12, 768, 64).unwrap();
        // Should fail with too many KV heads
        let invalid_result = base_config.clone().with_kv_heads(15); // More than num_heads
        assert!(invalid_result.is_err());
        assert!(invalid_result.unwrap_err().to_string().contains("cannot exceed"));
        // Should fail with non-divisible KV heads
        let invalid_result2 = base_config.clone().with_kv_heads(5); // 12 not divisible by 5
        assert!(invalid_result2.is_err());
        assert!(invalid_result2.unwrap_err().to_string().contains("divisible"));
        // Should succeed with valid KV heads
        let valid_result = base_config.with_kv_heads(6); // 12 divisible by 6
        assert!(valid_result.is_ok());
        let valid_config = valid_result.unwrap();
        assert_eq!(valid_config.get_num_kv_heads(), 6);
        assert!(valid_config.is_gqa());
    }
    #[test]
    fn test_sliding_window_window_size_effects_should_fail_initially() {
        // This test should FAIL initially - testing how window size affects computations
        let device = Device::cuda(0).unwrap_or(Device::default());
        let seq_len = 256;
        // Test different window sizes
        let small_window_config = SlidingWindowConfig::new(8, 512, 32).unwrap(); // Small window
        let large_window_config = SlidingWindowConfig::new(8, 512, 128).unwrap(); // Large window
        let small_attention = SlidingWindowAttention::new(small_window_config, &device).unwrap();
        let large_attention = SlidingWindowAttention::new(large_window_config, &device).unwrap();
        // Memory savings should be different
        let small_savings = small_attention.memory_savings_ratio(seq_len);
        let large_savings = large_attention.memory_savings_ratio(seq_len);
        // Smaller window should save more memory
        assert!(small_savings > large_savings);
        // Effective attention spans should be different
        let small_span = small_attention.effective_attention_span(seq_len);
        let large_span = large_attention.effective_attention_span(seq_len);
        assert!(large_span > small_span);
    }
    #[test]
    fn test_sliding_window_from_transformer_config_should_fail_initially() {
        // This test should FAIL initially - testing creation from TransformerConfig
        let mut transformer_config = TransformerConfig::new(50257, 768, 12, 12, 3072, 1024);
        // Test with standard MHA
        let result = SlidingWindowConfig::from_transformer_config(&transformer_config, 256);
        assert!(result.is_ok());
        let config = result.unwrap();
        assert_eq!(config.window_size, 256);
        assert_eq!(config.num_heads, 12);
        assert_eq!(config.d_model, 768);
        assert!(config.num_kv_heads.is_none()); // Standard MHA
        // Test with MQA enabled
        transformer_config.set_mqa(true, 1).unwrap();
        let mqa_result = SlidingWindowConfig::from_transformer_config(&transformer_config, 256);
        assert!(mqa_result.is_ok());
        let mqa_config = mqa_result.unwrap();
        assert_eq!(mqa_config.get_num_kv_heads(), 1);
        assert!(mqa_config.is_mqa());
        // Test with invalid window size
        let invalid_result = SlidingWindowConfig::from_transformer_config(&transformer_config, 0);
        assert!(invalid_result.is_err());
        assert!(invalid_result.unwrap_err().to_string().contains("Window size must be greater than 0"));
    }
    #[test]
    fn test_sliding_window_performance_characteristics_should_fail_initially() {
        // This test should FAIL initially - testing performance characteristics
        let config = SlidingWindowConfig::new(16, 1024, 128).unwrap();
        let device = Device::cuda(0).unwrap_or(Device::default());
        let attention = SlidingWindowAttention::new(config, &device).unwrap();
        // Test memory savings for various sequence lengths
        let test_sequences = vec![256, 512, 1024, 2048, 4096];
        for seq_len in test_sequences {
            let savings = attention.memory_savings_ratio(seq_len);
            let effective_span = attention.effective_attention_span(seq_len);
            // Memory savings should increase with sequence length
            if seq_len > 256 { // Window size + some buffer
                assert!(savings > 0.5, "Should save significant memory for seq_len {}", seq_len);
            }
            // Effective span should be bounded by window size for causal attention
            if seq_len > 128 { // If sequence is longer than window
                assert_eq!(effective_span, 129); // window_size + 1 for causal
            }
        }
    }
 }
@@ -1,62 +0,0 @@
 //! # RTX Transformers: Complete Transformer Training Infrastructure
 //!
 //! This crate provides a complete transformer training infrastructure that is 5-10x faster
 //! than PyTorch 2.5 while integrating revolutionary quantum, neuromorphic, and edge capabilities.
 //!
 //! ## Features
 //!
 //! - **Adam/AdamW Optimizers**: Industry-standard optimizers with superior performance
 //! - **Learning Rate Schedulers**: Warmup, cosine annealing, and advanced scheduling
 //! - **Layer Components**: LayerNorm, RMSNorm, positional encodings
 //! - **Transformer Architectures**: GPT, BERT, T5 with modern variants
 //! - **Revolutionary Integration**: Quantum-enhanced attention, neuromorphic preprocessing, edge deployment
 //!
 //! ## Examples
 //!
 //! ```rust
 //! use rtx_transformers::{AdamOptimizer, TransformerConfig};
 //!
 //! // Create Adam optimizer
 //! let optimizer = AdamOptimizer::new(0.001, 0.9, 0.999, 1e-8, 0.01)?;
 //!
 //! // Configure transformer
 //! let config = TransformerConfig::gpt2_small();
 //! # Ok::<(), Box<dyn std::error::Error>>(())
 //! ```
 #![deny(missing_docs)]
 #![deny(unsafe_code)]
 pub mod error;
 pub mod optimizers;
 pub mod schedulers;
 pub mod layers;
 pub mod architectures;
 pub mod training;
 pub mod tokenization;
 pub mod revolutionary;
 pub mod benchmarks;
 // Re-export key types for convenience
 pub use error::{TransformerError, Result};
 pub use optimizers::{AdamOptimizer, AdamWOptimizer, Optimizer};
 pub use schedulers::{LearningRateScheduler, WarmupScheduler, CosineAnnealingScheduler};
 pub use layers::{LayerNorm, RMSNorm, PositionalEncoding};
 pub use architectures::{TransformerConfig, TransformerBlock, GPTConfig};
 pub use training::{TrainingConfig, TrainingLoop};
 /// Current version of rtx-transformers
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 /// Verify transformer infrastructure is properly initialized
 pub fn verify_infrastructure() -> Result<()> {
    tracing::info!("RTX Transformers v{} - Verifying infrastructure", VERSION);
    // Verify core dependencies are available
    rtx_tensor::verify_backend()?;
    rtx_autograd::verify_gradient_system()?;
    rtx_runtime::verify_cuda_availability()?;
    tracing::info!("RTX Transformers infrastructure verification complete");
    Ok(())
 }
@@ -1,19 +0,0 @@
 //! # RTX Transformers: Complete Transformer Training Infrastructure (Minimal Version)
 //!
 //! This is a minimal version to test compilation with rustg tools.
 #![deny(missing_docs)]
 pub mod error;
 // Re-export key types for convenience
 pub use error::{TransformerError, Result};
 /// Current version of rtx-transformers
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 /// Simple test function to verify the crate loads
 pub fn test_minimal() -> Result<()> {
    tracing::info!("RTX Transformers v{} - Minimal test successful", VERSION);
    Ok(())
 }
@@ -1,380 +0,0 @@
 //! Batch 33B — Complete I-JEPA end-to-end forward pass integration.
 //!
 //! Wires together all JEPA components into a single `JepaFullPipeline`:
 //!  1. Block masking — sample context and target patch indices.
 //!  2. Context encoding — encode context patches with the online encoder.
 //!  3. Prediction — narrow predictor transformer maps context → target space.
 //!  4. Target encoding — encode target patches with the EMA encoder (no gradient).
 //!  5. Loss — L2 between predicted and target representations.
 use std::time::Instant;
 use super::jepa::{BlockMaskStrategy, JepaPredictor, JepaLossResult, jepa_loss};
 use super::jepa_vit::{JepaEncoder, JepaViTConfig, CpuViTEncoder, EmaTargetEncoderDyn};
 use super::jepa_data::{ImageRecord, JepaAugmentationPipeline, JepaDataConfig};
 // ============================================================================
 // JepaStepOutput
 // ============================================================================
 /// Output of one `JepaFullPipeline` step.
 #[derive(Debug, Clone)]
 pub struct JepaStepOutput {
    /// Per-block I-JEPA loss result.
    pub loss_result: JepaLossResult,
    /// Scalar loss (mean across target blocks).
    pub loss: f32,
    /// Context patch count.
    pub n_context: usize,
    /// Total target patch count.
    pub n_target: usize,
    /// Current EMA tau at the time of this step.
    pub ema_tau: f64,
    /// Wall-clock time for this step (µs).
    pub step_us: u64,
 }
 // ============================================================================
 // JepaFullPipeline
 // ============================================================================
 /// Complete I-JEPA training pipeline: context encoder + predictor + EMA target encoder.
 ///
 /// This is the reference implementation of the full I-JEPA forward pass:
 ///  1. Block masking: sample context and target patch indices.
 ///  2. Context encoding: encode context patches.
 ///  3. Prediction: predict target representations from context.
 ///  4. Target encoding: encode target patches with EMA encoder (no gradient).
 ///  5. Loss: L2 between predicted and target representations.
 pub struct JepaFullPipeline {
    /// Context (online) encoder.
    pub context_encoder: Box<dyn JepaEncoder>,
    /// Narrow predictor transformer (encoder_dim → predictor_dim → encoder_dim).
    pub predictor: JepaPredictor,
    /// EMA target encoder.
    pub target_encoder: EmaTargetEncoderDyn,
    /// Block masking strategy.
    pub mask_strategy: BlockMaskStrategy,
    /// Step counter.
    pub step: usize,
    /// EMA tau schedule: (tau_start, tau_end, total_steps).
    tau_schedule: (f64, f64, usize),
 }
 impl JepaFullPipeline {
    /// Create a new pipeline from a `JepaViTConfig` using a `CpuViTEncoder`.
    ///
    /// A second independent `CpuViTEncoder` (tiny preset) is used for the EMA
    /// target encoder.  Its weights diverge immediately via the tau schedule, so
    /// the asymmetry is expected at this stage.
    pub fn from_vit_config(
        vit_cfg: JepaViTConfig,
        predictor_dim: usize,
        predictor_depth: usize,
        tau_start: f64,
        tau_end: f64,
        total_steps: usize,
    ) -> Self {
        let embed_dim = vit_cfg.embed_dim;
        let num_patches = vit_cfg.num_patches();
        let num_heads = (predictor_dim / 64).max(1).min(8);
        let context_enc: Box<dyn JepaEncoder> = Box::new(CpuViTEncoder::new(vit_cfg));
        // Target encoder: independent CpuViTEncoder.
        // EmaTargetEncoderDyn wraps it and updates tau on each step.
        let target_enc_inner: Box<dyn JepaEncoder> =
            Box::new(CpuViTEncoder::new(JepaViTConfig::tiny()));
        let target_enc = EmaTargetEncoderDyn::from_encoder(target_enc_inner, tau_start);
        let predictor = JepaPredictor::new(
            embed_dim,
            predictor_dim,
            predictor_depth,
            num_heads,
            num_patches,
        );
        Self {
            context_encoder: context_enc,
            predictor,
            target_encoder: target_enc,
            mask_strategy: BlockMaskStrategy::default_ijepa(),
            step: 0,
            tau_schedule: (tau_start, tau_end, total_steps),
        }
    }
    /// Create with a custom pair of encoders (e.g. `GpuViTEncoder`).
    ///
    /// `context_enc` is used as the online encoder; `target_enc_clone` is wrapped
    /// in `EmaTargetEncoderDyn` as the EMA target encoder.
    pub fn with_encoder(
        context_enc: Box<dyn JepaEncoder>,
        target_enc_clone: Box<dyn JepaEncoder>,
        predictor_dim: usize,
        predictor_depth: usize,
        tau_start: f64,
        tau_end: f64,
        total_steps: usize,
    ) -> Self {
        let embed_dim = context_enc.embed_dim();
        let num_patches = context_enc.num_patches();
        let num_heads = (predictor_dim / 64).max(1).min(8);
        let predictor =
            JepaPredictor::new(embed_dim, predictor_dim, predictor_depth, num_heads, num_patches);
        let target_enc = EmaTargetEncoderDyn::from_encoder(target_enc_clone, tau_start);
        Self {
            context_encoder: context_enc,
            predictor,
            target_encoder: target_enc,
            mask_strategy: BlockMaskStrategy::default_ijepa(),
            step: 0,
            tau_schedule: (tau_start, tau_end, total_steps),
        }
    }
    /// Current EMA tau derived from the linear schedule.
    pub fn current_tau(&self) -> f64 {
        let (start, end, total) = self.tau_schedule;
        if total == 0 {
            return end;
        }
        start + (end - start) * (self.step as f64 / total as f64).min(1.0)
    }
    /// Run one training step with externally-provided patch indices.
    ///
    /// Core I-JEPA forward pass:
    ///   context_encoder(context_patches) → predictor → compare with target_encoder(target_patches)
    pub fn step_with_indices(
        &mut self,
        context_indices: &[usize],
        target_indices: &[usize],
    ) -> JepaStepOutput {
        let t0 = Instant::now();
        self.step += 1;
        let tau = self.current_tau();
        // 1. Context encoding
        let ctx_reps = self.context_encoder.encode(context_indices);
        // 2. Target encoding (EMA, no gradient)
        let tgt_reps = self.target_encoder.encode(target_indices);
        // 3. Prediction: context representations → predicted target representations
        let predictions = self.predictor.forward(&ctx_reps, context_indices, target_indices);
        let encoder_dim = self.context_encoder.embed_dim();
        // 4. Loss: predicted vs target (L2 in representation space)
        let target_blocks = vec![target_indices.to_vec()];
        let target_block_offsets = vec![0usize];
        let loss_result = jepa_loss(
            &predictions,
            &tgt_reps,
            encoder_dim,
            &target_blocks,
            &target_block_offsets,
        );
        // 5. Update EMA tau
        self.target_encoder.update_tau(tau);
        let step_us = t0.elapsed().as_micros() as u64;
        JepaStepOutput {
            loss: loss_result.loss,
            loss_result,
            n_context: context_indices.len(),
            n_target: target_indices.len(),
            ema_tau: tau,
            step_us,
        }
    }
    /// Run one training step using block masking (masks generated internally).
    ///
    /// `image_seed` deterministically varies the mask per sample.
    pub fn step_with_mask(&mut self, image_seed: u64) -> JepaStepOutput {
        let num_patches = self.context_encoder.num_patches();
        // Approximate a square grid; non-square images are rare in JEPA pre-training.
        let grid = (num_patches as f64).sqrt() as usize;
        let grid_h = grid;
        let grid_w = (num_patches + grid - 1) / grid;
        let mask = self.mask_strategy.generate(grid_h, grid_w, image_seed);
        self.step_with_indices(&mask.context_indices, &mask.all_target_indices)
    }
    /// Run one step with a real image: augment → mask → step.
    ///
    /// The augmented pixel buffer is computed but not yet fed into the encoder
    /// (CpuViTEncoder generates its own deterministic embeddings from patch indices).
    /// A full implementation would pass `pixels` through a patch-embedding layer.
    pub fn step_with_image(
        &mut self,
        image: &ImageRecord,
        aug_pipeline: &JepaAugmentationPipeline,
        seed: u64,
    ) -> JepaStepOutput {
        // Augment the image (side-effectfully exercises the data pipeline).
        let _pixels = aug_pipeline.process(image, seed);
        // Delegate to mask-based step (pixel values not consumed by CpuViTEncoder).
        self.step_with_mask(seed)
    }
 }
 // ============================================================================
 // Tests
 // ============================================================================
 #[cfg(test)]
 mod tests {
    use super::*;
    use crate::ssl::jepa_vit::JepaViTConfig;
    use crate::ssl::jepa_data::{ImageRecord, JepaAugmentationPipeline, JepaDataConfig};
    /// A tiny pipeline for fast tests: embed_dim=32, depth=1, 2 heads.
    fn tiny_pipeline() -> JepaFullPipeline {
        let vit_cfg = JepaViTConfig {
            embed_dim: 32,
            depth: 1,
            num_heads: 2,
            mlp_ratio: 2.0,
            patch_size: 16,
            image_size: 64,
        };
        JepaFullPipeline::from_vit_config(vit_cfg, 16, 1, 0.996, 1.0, 100)
    }
    // 1. step_with_indices returns finite loss
    #[test]
    fn test_step_with_indices_finite() {
        let mut p = tiny_pipeline();
        let out = p.step_with_indices(&[0, 1, 2, 3, 4, 5], &[6, 7, 8, 9]);
        assert!(out.loss.is_finite(), "loss={}", out.loss);
    }
    // 2. n_context matches input
    #[test]
    fn test_step_context_count() {
        let mut p = tiny_pipeline();
        let out = p.step_with_indices(&[0, 1, 2], &[4, 5]);
        assert_eq!(out.n_context, 3);
    }
    // 3. n_target matches input
    #[test]
    fn test_step_target_count() {
        let mut p = tiny_pipeline();
        let out = p.step_with_indices(&[0, 1, 2], &[4, 5, 6, 7]);
        assert_eq!(out.n_target, 4);
    }
    // 4. step counter increments
    #[test]
    fn test_step_counter() {
        let mut p = tiny_pipeline();
        p.step_with_indices(&[0], &[1]);
        p.step_with_indices(&[0], &[1]);
        assert_eq!(p.step, 2);
    }
    // 5. step_with_mask produces finite loss
    #[test]
    fn test_step_with_mask_finite() {
        let mut p = tiny_pipeline();
        let out = p.step_with_mask(42);
        assert!(out.loss.is_finite(), "mask-step loss={}", out.loss);
    }
    // 6. step_with_mask: mask-generated n_context > 0
    #[test]
    fn test_step_with_mask_context_nonzero() {
        let mut p = tiny_pipeline();
        let out = p.step_with_mask(0);
        assert!(out.n_context > 0, "must have context patches");
    }
    // 7. step_with_mask: mask-generated n_target > 0
    #[test]
    fn test_step_with_mask_target_nonzero() {
        let mut p = tiny_pipeline();
        let out = p.step_with_mask(0);
        assert!(out.n_target > 0, "must have target patches");
    }
    // 8. current_tau at step 0 = tau_start
    #[test]
    fn test_tau_at_step_zero() {
        let p = tiny_pipeline();
        let tau = p.current_tau();
        assert!((tau - 0.996).abs() < 1e-6, "tau={tau}");
    }
    // 9. tau increases toward tau_end over steps
    #[test]
    fn test_tau_increases() {
        let mut p = tiny_pipeline();
        let tau0 = p.current_tau();
        p.step_with_mask(0);
        p.step_with_mask(1);
        let tau1 = p.current_tau();
        assert!(tau1 >= tau0, "tau must not decrease");
    }
    // 10. step_with_image runs without panic
    #[test]
    fn test_step_with_image() {
        let mut p = tiny_pipeline();
        let image = ImageRecord {
            pixels: vec![0.5f32; 64 * 64 * 3],
            width: 64,
            height: 64,
            channels: 3,
            label: None,
            key: "test".to_string(),
        };
        // Use the actual JepaDataConfig fields (not those in the spec).
        let cfg = JepaDataConfig {
            image_size: 64,
            patch_size: 16,
            batch_size: 1,
            num_workers: 0,
            shard_paths: Vec::new(),
            scale_range: (0.2, 1.0),
            ratio_range: (0.75, 1.33),
            use_horizontal_flip: false,
            imagenet_normalize: false,
        };
        let aug = JepaAugmentationPipeline::from_config(&cfg);
        let out = p.step_with_image(&image, &aug, 42);
        assert!(out.loss.is_finite());
    }
    // 11. step_us is a valid u64 (timing works)
    #[test]
    fn test_step_timing() {
        let mut p = tiny_pipeline();
        let out = p.step_with_mask(0);
        // step_us can be 0 on very fast systems; just assert it is a valid value.
        let _ = out.step_us;
    }
    // 12. JepaStepOutput.loss_result.loss == JepaStepOutput.loss
    #[test]
    fn test_loss_consistency() {
        let mut p = tiny_pipeline();
        let out = p.step_with_indices(&[0, 1, 2, 3], &[5, 6, 7, 8]);
        assert!(
            (out.loss - out.loss_result.loss).abs() < 1e-6,
            "loss={} vs loss_result.loss={}",
            out.loss,
            out.loss_result.loss
        );
    }
 }
@@ -1,647 +0,0 @@
 //! SimMIM (Simple Masked Image Modeling) Implementation
 //!
 //! SimMIM is a simple framework for masked image modeling that directly predicts 
 //! raw pixel values of masked patches. Based on "SimMIM: A Simple Framework for 
 //! Masked Image Modeling" (Xie et al., 2021).
 //!
 //! Key features:
 //! - Random masking of image patches (60% for ViT, 32% for Swin)
 //! - Direct raw pixel regression (no tokenizer needed)
 //! - Simple prediction head (linear layer)  
 //! - L1 loss for pixel prediction
 //! - Support for both ViT and Swin Transformer
 //! 
 //! Algorithm:
 //! 1. Random masking: Randomly mask patches
 //! 2. Mask tokens: Replace masked patches with learnable mask token
 //! 3. Encoder: Pass through ViT/Swin transformer
 //! 4. Prediction: Simple linear layer to predict raw pixels
 //! 5. Loss: L1 loss between predicted and actual pixels (only on masked patches)
 use crate::prelude::*;
 use std::sync::Arc;
 use parking_lot::RwLock;
 /// Configuration for different encoder types
 #[derive(Debug, Clone, PartialEq)]
 pub enum EncoderType {
    /// Vision Transformer encoder
    ViT,
    /// Swin Transformer encoder  
    Swin,
 }
 /// Type of prediction head
 #[derive(Debug, Clone, PartialEq)]
 pub enum PredictionHeadType {
    /// Simple linear layer
    Linear,
 }
 /// Loss function type
 #[derive(Debug, Clone, PartialEq)]
 pub enum LossType {
    /// L1 (MAE) loss
    L1,
    /// L2 (MSE) loss
    L2,
 }
 /// SimMIM configuration parameters
 #[derive(Debug, Clone)]
 pub struct SimMIMConfig {
    /// Fraction of patches to mask
    pub mask_ratio: f32,
    /// Mask patch size (for hierarchical masking in Swin)
    pub mask_patch_size: usize,
    /// Type of prediction head
    pub prediction_head: PredictionHeadType,
    /// Loss function type
    pub loss_type: LossType,
    /// Whether to normalize pixel targets
    pub norm_pix_loss: bool,
    /// Encoder embedding dimension
    pub encoder_dim: usize,
    /// Image patch size
    pub patch_size: usize,
    /// Type of encoder (ViT or Swin)
    pub encoder_type: EncoderType,
    /// Input image size
    pub image_size: usize,
    /// Number of input channels
    pub in_channels: usize,
 }
 impl Default for SimMIMConfig {
    fn default() -> Self {
        Self {
            mask_ratio: 0.6,
            mask_patch_size: 32,
            prediction_head: PredictionHeadType::Linear,
            loss_type: LossType::L1,
            norm_pix_loss: false,
            encoder_dim: 768,
            patch_size: 16,
            encoder_type: EncoderType::ViT,
            image_size: 224,
            in_channels: 3,
        }
    }
 }
 impl SimMIMConfig {
    /// Create configuration optimized for ViT
    pub fn for_vit() -> Self {
        Self {
            encoder_type: EncoderType::ViT,
            mask_ratio: 0.6,
            mask_patch_size: 16,
            ..Default::default()
        }
    }
    /// Create configuration optimized for Swin Transformer
    pub fn for_swin() -> Self {
        Self {
            encoder_type: EncoderType::Swin,
            mask_ratio: 0.32, // Lower mask ratio for hierarchical architecture
            mask_patch_size: 32, // Larger patch size for hierarchical masking
            ..Default::default()
        }
    }
    /// Set mask ratio
    pub fn with_mask_ratio(mut self, mask_ratio: f32) -> Self {
        self.mask_ratio = mask_ratio.clamp(0.0, 1.0);
        self
    }
    /// Set encoder type
    pub fn with_encoder_type(mut self, encoder_type: EncoderType) -> Self {
        self.encoder_type = encoder_type;
        self
    }
    /// Set whether to normalize pixel loss
    pub fn with_norm_pix_loss(mut self, norm_pix_loss: bool) -> Self {
        self.norm_pix_loss = norm_pix_loss;
        self
    }
    /// Set encoder dimension
    pub fn with_encoder_dim(mut self, encoder_dim: usize) -> Self {
        self.encoder_dim = encoder_dim;
        self
    }
 }
 /// Result from random patch masking
 #[derive(Debug)]
 pub struct PatchMaskResult {
    /// Boolean mask for each patch (true = masked, false = visible)
    pub mask_indices: Vec<Vec<bool>>,
    /// Indices of visible patches for each batch item
    pub visible_indices: Vec<Vec<usize>>,
    /// Total number of masked patches per batch item
    pub num_masked_patches: usize,
 }
 /// Random patch masker for SimMIM
 #[derive(Debug)]
 pub struct RandomPatchMasker {
    mask_ratio: f32,
 }
 impl RandomPatchMasker {
    /// Create new random patch masker
    pub fn new(mask_ratio: f32) -> Self {
        Self {
            mask_ratio: mask_ratio.clamp(0.0, 1.0),
        }
    }
    /// Get mask ratio
    pub fn mask_ratio(&self) -> f32 {
        self.mask_ratio
    }
    /// Generate random patch mask for a batch
    pub fn mask_patches(&self, patches: &Tensor, seed: Option<u64>) -> Result<PatchMaskResult> {
        let shape = patches.shape();
        if shape.len() != 3 {
            return Err(TransformerError::InvalidInput("Patches tensor must be 3D [batch, patches, embed_dim]".into()));
        }
        let batch_size = shape[0];
        let num_patches = shape[1];
        if num_patches == 0 {
            return Err(TransformerError::InvalidInput("Number of patches cannot be zero".into()));
        }
        let num_masked = (num_patches as f32 * self.mask_ratio).round() as usize;
        let num_visible = num_patches - num_masked;
        let mut mask_indices = Vec::new();
        let mut visible_indices = Vec::new();
        let mut rng_state = seed.unwrap_or(42);
        for _ in 0..batch_size {
            // Create all patch indices
            let mut indices: Vec<usize> = (0..num_patches).collect();
            // Shuffle indices deterministically
            for i in (1..indices.len()).rev() {
                rng_state = rng_state.wrapping_mul(1103515245).wrapping_add(12345);
                let j = (rng_state as usize) % (i + 1);
                indices.swap(i, j);
            }
            // Select visible patches
            let visible = indices[..num_visible].to_vec();
            visible_indices.push(visible.clone());
            // Create boolean mask
            let mut mask = vec![true; num_patches]; // Start with all masked
            for &idx in &visible {
                mask[idx] = false; // Mark visible patches
            }
            mask_indices.push(mask);
        }
        Ok(PatchMaskResult {
            mask_indices,
            visible_indices,
            num_masked_patches: num_masked,
        })
    }
 }
 /// Learnable mask token embedding
 #[derive(Debug)]
 pub struct MaskTokenEmbedding {
    token: Arc<RwLock<Tensor>>,
    embed_dim: usize,
    device: Device,
 }
 impl MaskTokenEmbedding {
    /// Create new mask token embedding
    pub fn new(embed_dim: usize, device: &Device) -> Result<Self> {
        let token = Tensor::randn(vec![1, embed_dim], DType::F32, device)?;
        Ok(Self {
            token: Arc::new(RwLock::new(token)),
            embed_dim,
            device: device.clone(),
        })
    }
    /// Get embedding dimension
    pub fn embed_dim(&self) -> usize {
        self.embed_dim
    }
    /// Get the mask token
    pub fn get_token(&self) -> Result<Tensor> {
        let token = self.token.read();
        Ok(token.clone())
    }
    /// Broadcast mask token to specified shape
    pub fn broadcast(&self, batch_size: usize, num_masked: usize) -> Result<Tensor> {
        let token = self.token.read();
        // Create broadcasted tensor
        let token_data = token.to_vec::<f32>()?;
        let mut broadcasted_data = vec![0.0f32; batch_size * num_masked * self.embed_dim];
        for b in 0..batch_size {
            for m in 0..num_masked {
                let base_idx = b * num_masked * self.embed_dim + m * self.embed_dim;
                for d in 0..self.embed_dim {
                    broadcasted_data[base_idx + d] = token_data[d];
                }
            }
        }
        Tensor::from_data(
            broadcasted_data,
            vec![batch_size, num_masked, self.embed_dim],
            DType::F32,
            &self.device,
        )
    }
 }
 /// Linear prediction head for pixel reconstruction
 #[derive(Debug)]
 pub struct LinearPredictionHead {
    linear: Arc<RwLock<Tensor>>,
    bias: Arc<RwLock<Tensor>>,
    input_dim: usize,
    output_dim: usize,
    device: Device,
 }
 impl LinearPredictionHead {
    /// Create new linear prediction head
    pub fn new(input_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
        let linear = Tensor::randn(vec![input_dim, output_dim], DType::F32, device)?;
        let bias = Tensor::zeros(vec![output_dim], device)?;
        Ok(Self {
            linear: Arc::new(RwLock::new(linear)),
            bias: Arc::new(RwLock::new(bias)),
            input_dim,
            output_dim,
            device: device.clone(),
        })
    }
    /// Get input dimension
    pub fn input_dim(&self) -> usize {
        self.input_dim
    }
    /// Get output dimension
    pub fn output_dim(&self) -> usize {
        self.output_dim
    }
    /// Forward pass through prediction head
    pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
        let linear = self.linear.read();
        let bias = self.bias.read();
        let output = input.matmul(&*linear)?;
        output.add(&*bias)
    }
 }
 /// SimMIM loss computation utilities
 pub struct SimMIMLoss;
 impl SimMIMLoss {
    /// Compute L1 (MAE) loss between predicted and target pixels
    pub fn compute_l1_loss(predicted: &Tensor, target: &Tensor) -> Result<Tensor> {
        // Check tensor shapes match
        let pred_shape = predicted.shape();
        let target_shape = target.shape();
        if pred_shape != target_shape {
            return Err(TransformerError::InvalidInput(
                format!("Predicted and target shapes must match: {:?} vs {:?}", pred_shape, target_shape)
            ));
        }
        // Compute L1 loss: mean(|predicted - target|)
        let diff = predicted.sub(target)?;
        let abs_diff = diff.abs()?;
        abs_diff.mean(&[])
    }
    /// Compute normalized L1 loss (normalize pixel targets first)
    pub fn compute_l1_loss_normalized(predicted: &Tensor, target: &Tensor) -> Result<Tensor> {
        // Normalize target pixels to have zero mean and unit variance per patch
        let normalized_target = Self::normalize_pixels(target)?;
        let normalized_predicted = Self::normalize_pixels(predicted)?;
        Self::compute_l1_loss(&normalized_predicted, &normalized_target)
    }
    /// Compute L2 (MSE) loss between predicted and target pixels  
    pub fn compute_l2_loss(predicted: &Tensor, target: &Tensor) -> Result<Tensor> {
        let pred_shape = predicted.shape();
        let target_shape = target.shape();
        if pred_shape != target_shape {
            return Err(TransformerError::InvalidInput(
                format!("Predicted and target shapes must match: {:?} vs {:?}", pred_shape, target_shape)
            ));
        }
        // Compute L2 loss: mean((predicted - target)^2)
        let diff = predicted.sub(target)?;
        let squared_diff = diff.pow_scalar(2.0)?;
        squared_diff.mean(&[])
    }
    fn normalize_pixels(pixels: &Tensor) -> Result<Tensor> {
        // Simple normalization: (x - mean) / std per patch
        let mean = pixels.mean(&[2], true)?; // Keep last dimension for broadcasting
        let centered = pixels.sub(&mean)?;
        let var = centered.pow_scalar(2.0)?.mean(&[2], true)?;
        let std = var.sqrt()?;
        // Add small epsilon to avoid division by zero
        let eps = Tensor::full(&std.shape(), 1e-6, DType::F32, pixels.device())?;
        let std_safe = std.add(&eps)?;
        centered.div(&std_safe)
    }
 }
 /// Result from SimMIM training step
 #[derive(Debug)]
 pub struct SimMIMTrainingResult {
    /// Reconstruction loss
    pub loss: Tensor,
    /// Number of masked patches in this batch
    pub num_masked_patches: usize,
    /// Actual mask ratio used
    pub mask_ratio: f32,
 }
 /// Main SimMIM trainer
 #[derive(Debug)]
 pub struct SimMIMTrainer {
    config: SimMIMConfig,
    masker: RandomPatchMasker,
    mask_token: MaskTokenEmbedding,
    prediction_head: LinearPredictionHead,
    encoder: Arc<RwLock<Tensor>>, // Simplified encoder (in practice would use full ViT/Swin)
    device: Device,
    is_training: bool,
 }
 impl SimMIMTrainer {
    /// Create new SimMIM trainer
    pub fn new(
        config: SimMIMConfig,
        in_channels: usize,
        image_size: usize,
        device: &Device,
    ) -> Result<Self> {
        // Create masker
        let masker = RandomPatchMasker::new(config.mask_ratio);
        // Create mask token embedding
        let mask_token = MaskTokenEmbedding::new(config.encoder_dim, device)?;
        // Create prediction head
        let patch_volume = config.patch_size * config.patch_size * in_channels;
        let prediction_head = LinearPredictionHead::new(config.encoder_dim, patch_volume, device)?;
        // Create simplified encoder (in practice would use full transformer)
        let num_patches = (image_size / config.patch_size).pow(2);
        let encoder = Tensor::randn(
            vec![patch_volume, config.encoder_dim],
            DType::F32,
            device,
        )?;
        Ok(Self {
            config,
            masker,
            mask_token,
            prediction_head,
            encoder: Arc::new(RwLock::new(encoder)),
            device: device.clone(),
            is_training: true,
        })
    }
    /// Get configuration
    pub fn config(&self) -> &SimMIMConfig {
        &self.config
    }
    /// Get device
    pub fn device(&self) -> &Device {
        &self.device
    }
    /// Set to training mode
    pub fn train(&mut self) {
        self.is_training = true;
    }
    /// Set to evaluation mode
    pub fn eval(&mut self) {
        self.is_training = false;
    }
    /// Check if in training mode
    pub fn is_training(&self) -> bool {
        self.is_training
    }
    /// Perform one SimMIM training step
    pub fn train_step(&mut self, images: &Tensor, seed: Option<u64>) -> Result<SimMIMTrainingResult> {
        let batch_size = images.shape()[0];
        // Step 1: Convert images to patches
        let patches = self.extract_patches(images)?;
        let target_pixels = patches.clone(); // Keep original patches as targets
        // Step 2: Generate random mask
        let mask_result = self.masker.mask_patches(&patches, seed)?;
        // Step 3: Apply masking and encode
        let masked_patches = self.apply_mask_tokens(&patches, &mask_result)?;
        let encoded_features = self.encode_patches(&masked_patches)?;
        // Step 4: Predict pixels for masked patches only
        let masked_features = self.extract_masked_features(&encoded_features, &mask_result)?;
        let predicted_pixels = self.prediction_head.forward(&masked_features)?;
        // Step 5: Extract target pixels for masked patches
        let target_masked_pixels = self.extract_masked_targets(&target_pixels, &mask_result)?;
        // Step 6: Compute loss (L1 loss on masked patches only)
        let loss = match self.config.loss_type {
            LossType::L1 => {
                if self.config.norm_pix_loss {
                    SimMIMLoss::compute_l1_loss_normalized(&predicted_pixels, &target_masked_pixels)?
                } else {
                    SimMIMLoss::compute_l1_loss(&predicted_pixels, &target_masked_pixels)?
                }
            }
            LossType::L2 => {
                SimMIMLoss::compute_l2_loss(&predicted_pixels, &target_masked_pixels)?
            }
        };
        Ok(SimMIMTrainingResult {
            loss,
            num_masked_patches: mask_result.num_masked_patches,
            mask_ratio: self.masker.mask_ratio(),
        })
    }
    /// Extract features in evaluation mode (no masking)
    pub fn extract_features(&self, images: &Tensor) -> Result<Tensor> {
        let patches = self.extract_patches(images)?;
        let encoded = self.encode_patches(&patches)?;
        // Global average pooling to get image-level features
        encoded.mean(&[1]) // Average over patch dimension
    }
    fn extract_patches(&self, images: &Tensor) -> Result<Tensor> {
        // Simplified patch extraction - reshape image into patches
        let shape = images.shape();
        let batch_size = shape[0];
        let channels = shape[1];
        let height = shape[2];
        let width = shape[3];
        let patch_size = self.config.patch_size;
        let patches_per_row = height / patch_size;
        let patches_per_col = width / patch_size;
        let num_patches = patches_per_row * patches_per_col;
        let patch_volume = patch_size * patch_size * channels;
        // Simulate patch extraction by reshaping
        images.reshape(&[batch_size, num_patches, patch_volume])
    }
    fn apply_mask_tokens(&self, patches: &Tensor, mask_result: &PatchMaskResult) -> Result<Tensor> {
        let shape = patches.shape();
        let batch_size = shape[0];
        let num_patches = shape[1];
        let patch_volume = shape[2];
        // Create tensor with mask tokens replacing masked patches
        let mut masked_data = vec![0.0f32; batch_size * num_patches * patch_volume];
        let patches_data = patches.to_vec::<f32>()?;
        // Get mask token data
        let mask_token = self.mask_token.get_token()?;
        let mask_token_data = mask_token.to_vec::<f32>()?;
        for b in 0..batch_size {
            for p in 0..num_patches {
                let base_idx = b * num_patches * patch_volume + p * patch_volume;
                if mask_result.mask_indices[b][p] {
                    // Use mask token (broadcast to patch volume)
                    for i in 0..patch_volume {
                        masked_data[base_idx + i] = mask_token_data[i % self.config.encoder_dim];
                    }
                } else {
                    // Keep original patch
                    for i in 0..patch_volume {
                        masked_data[base_idx + i] = patches_data[base_idx + i];
                    }
                }
            }
        }
        Tensor::from_data(masked_data, shape, &device)
    }
    fn encode_patches(&self, patches: &Tensor) -> Result<Tensor> {
        let encoder = self.encoder.read();
        patches.matmul(&*encoder)
    }
    fn extract_masked_features(&self, features: &Tensor, mask_result: &PatchMaskResult) -> Result<Tensor> {
        let shape = features.shape();
        let batch_size = shape[0];
        let num_patches = shape[1];
        let embed_dim = shape[2];
        let num_masked = mask_result.num_masked_patches;
        let mut masked_features = vec![0.0f32; batch_size * num_masked * embed_dim];
        let features_data = features.to_vec::<f32>()?;
        for b in 0..batch_size {
            let mut masked_idx = 0;
            for p in 0..num_patches {
                if mask_result.mask_indices[b][p] && masked_idx < num_masked {
                    let src_base = b * num_patches * embed_dim + p * embed_dim;
                    let dst_base = b * num_masked * embed_dim + masked_idx * embed_dim;
                    for d in 0..embed_dim {
                        masked_features[dst_base + d] = features_data[src_base + d];
                    }
                    masked_idx += 1;
                }
            }
        }
        Tensor::from_data(
            masked_features,
            vec![batch_size, num_masked, embed_dim],
            DType::F32,
            &self.device,
        )
    }
    fn extract_masked_targets(&self, targets: &Tensor, mask_result: &PatchMaskResult) -> Result<Tensor> {
        let shape = targets.shape();
        let batch_size = shape[0];
        let num_patches = shape[1];
        let patch_volume = shape[2];
        let num_masked = mask_result.num_masked_patches;
        let mut masked_targets = vec![0.0f32; batch_size * num_masked * patch_volume];
        let targets_data = targets.to_vec::<f32>()?;
        for b in 0..batch_size {
            let mut masked_idx = 0;
            for p in 0..num_patches {
                if mask_result.mask_indices[b][p] && masked_idx < num_masked {
                    let src_base = b * num_patches * patch_volume + p * patch_volume;
                    let dst_base = b * num_masked * patch_volume + masked_idx * patch_volume;
                    for d in 0..patch_volume {
                        masked_targets[dst_base + d] = targets_data[src_base + d];
                    }
                    masked_idx += 1;
                }
            }
        }
        Tensor::from_data(
            masked_targets,
            vec![batch_size, num_masked, patch_volume],
            DType::F32,
            &self.device,
        )
    }
 }
@@ -1,39 +0,0 @@
 //! Simple standalone test for SimMIM compilation
 //! This test verifies the basic structure works before running the full test suite
 use crate::prelude::*;
 #[test]
 fn test_simmim_basic_compilation() {
    // Just test that types can be created
    use super::simmim::*;
    let mask_ratio = 0.6;
    let _config = SimMIMConfig::default().with_mask_ratio(mask_ratio);
    // Test basic masker creation
    let _masker = RandomPatchMasker::new(mask_ratio);
    assert!(true, "SimMIM types compile successfully");
 }
 #[test]
 fn test_simmim_config_methods() {
    use super::simmim::*;
    let config = SimMIMConfig::default()
        .with_mask_ratio(0.75)
        .with_encoder_type(EncoderType::Swin);
    assert_eq!(config.mask_ratio, 0.75);
    assert_eq!(config.encoder_type, EncoderType::Swin);
 }
 #[test] 
 fn test_simmim_enums() {
    use super::simmim::*;
    assert_eq!(EncoderType::ViT, EncoderType::ViT);
    assert_eq!(PredictionHeadType::Linear, PredictionHeadType::Linear);
    assert_eq!(LossType::L1, LossType::L1);
 }
@@ -1,522 +0,0 @@
 //! Comprehensive test suite for SimMIM (Simple Masked Image Modeling) implementation
 //! 
 //! Tests follow strict TDD - these failing tests drive the implementation.
 //! All tests must pass without mocks, stubs, or TODOs.
 use crate::prelude::*;
 use super::simmim::*;
 /// Test random patch masking functionality
 #[cfg(all(test, feature = "disabled_tests"))]
 mod random_masking_tests {
    use super::*;
    #[test]
    fn test_random_masker_creation() {
        let mask_ratio = 0.6;
        let masker = RandomPatchMasker::new(mask_ratio);
        assert_eq!(masker.mask_ratio(), mask_ratio);
    }
    #[test]
    fn test_random_masker_mask_ratio_validation() {
        // Test valid mask ratios
        let valid_ratios = [0.0, 0.3, 0.6, 0.75, 1.0];
        for ratio in valid_ratios {
            let masker = RandomPatchMasker::new(ratio);
            assert!((masker.mask_ratio() - ratio).abs() < 1e-6);
        }
    }
    #[test]
    fn test_random_patch_masking() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let masker = RandomPatchMasker::new(0.6);
        let batch_size = 2;
        let num_patches = 196; // 14x14 patches
        let embed_dim = 768;
        let patches = Tensor::randn(vec![batch_size, num_patches, embed_dim], DType::F32, &device).unwrap();
        let result = masker.mask_patches(&patches, Some(42)).unwrap();
        // Check mask structure
        assert_eq!(result.mask_indices.len(), batch_size);
        for mask in &result.mask_indices {
            assert_eq!(mask.len(), num_patches);
            let masked_count = mask.iter().filter(|&&x| x).count();
            let visible_count = mask.iter().filter(|&&x| !x).count();
            // Should mask approximately 60% of patches
            let expected_masked = (num_patches as f32 * 0.6).round() as usize;
            assert!((masked_count as i32 - expected_masked as i32).abs() <= 2);
            assert_eq!(masked_count + visible_count, num_patches);
        }
        // Check visible indices consistency
        assert_eq!(result.visible_indices.len(), batch_size);
        for (b, indices) in result.visible_indices.iter().enumerate() {
            for &idx in indices {
                assert!(!result.mask_indices[b][idx]); // visible indices should not be masked
            }
        }
    }
    #[test]
    fn test_random_masking_deterministic_with_seed() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let masker = RandomPatchMasker::new(0.6);
        let patches = Tensor::randn(vec![1, 196, 768], DType::F32, &device).unwrap();
        // Same seed should produce same mask
        let result1 = masker.mask_patches(&patches, Some(42)).unwrap();
        let result2 = masker.mask_patches(&patches, Some(42)).unwrap();
        assert_eq!(result1.mask_indices, result2.mask_indices);
        assert_eq!(result1.visible_indices, result2.visible_indices);
    }
    #[test]
    fn test_different_mask_ratios() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let patches = Tensor::randn(vec![1, 196, 768], DType::F32, &device).unwrap();
        let ratios = [0.3, 0.6, 0.75];
        for ratio in ratios {
            let masker = RandomPatchMasker::new(ratio);
            let result = masker.mask_patches(&patches, Some(42)).unwrap();
            let masked_count = result.mask_indices[0].iter().filter(|&&x| x).count();
            let expected = (196.0 * ratio).round() as usize;
            // Allow small deviation due to rounding
            assert!((masked_count as i32 - expected as i32).abs() <= 2);
        }
    }
 }
 /// Test mask token embedding functionality  
 #[cfg(all(test, feature = "disabled_tests"))]
 mod mask_token_tests {
    use super::*;
    #[test]
    fn test_mask_token_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let embed_dim = 768;
        let mask_token = MaskTokenEmbedding::new(embed_dim, &device).unwrap();
        assert_eq!(mask_token.embed_dim(), embed_dim);
    }
    #[test]
    fn test_mask_token_shape() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let embed_dim = 512;
        let mask_token = MaskTokenEmbedding::new(embed_dim, &device).unwrap();
        let token = mask_token.get_token().unwrap();
        assert_eq!(token.shape(), vec![1, embed_dim]);
    }
    #[test]
    fn test_mask_token_broadcast() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let embed_dim = 768;
        let batch_size = 4;
        let num_masked = 100;
        let mask_token = MaskTokenEmbedding::new(embed_dim, &device).unwrap();
        let broadcasted = mask_token.broadcast(batch_size, num_masked).unwrap();
        assert_eq!(broadcasted.shape(), vec![batch_size, num_masked, embed_dim]);
    }
 }
 /// Test linear prediction head functionality
 #[cfg(all(test, feature = "disabled_tests"))]
 mod prediction_head_tests {
    use super::*;
    #[test]
    fn test_linear_prediction_head_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let embed_dim = 768;
        let patch_size = 16;
        let in_channels = 3;
        let output_dim = patch_size * patch_size * in_channels; // 768
        let head = LinearPredictionHead::new(embed_dim, output_dim, &device).unwrap();
        assert_eq!(head.input_dim(), embed_dim);
        assert_eq!(head.output_dim(), output_dim);
    }
    #[test]
    fn test_prediction_head_forward() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let embed_dim = 768;
        let output_dim = 768; // 16x16x3
        let batch_size = 2;
        let num_masked = 100;
        let head = LinearPredictionHead::new(embed_dim, output_dim, &device).unwrap();
        let input = Tensor::randn(vec![batch_size, num_masked, embed_dim], DType::F32, &device).unwrap();
        let output = head.forward(&input).unwrap();
        assert_eq!(output.shape(), vec![batch_size, num_masked, output_dim]);
    }
    #[test]
    fn test_prediction_head_different_dimensions() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        // Test various common configurations
        let configs = vec![
            (768, 768),   // ViT-Base 16x16 patch
            (1024, 1024), // ViT-Large 16x16 patch  
            (512, 3072),  // Custom config
        ];
        for (input_dim, output_dim) in configs {
            let head = LinearPredictionHead::new(input_dim, output_dim, &device).unwrap();
            let input = Tensor::randn(vec![1, 10, input_dim], DType::F32, &device).unwrap();
            let output = head.forward(&input).unwrap();
            assert_eq!(output.shape(), vec![1, 10, output_dim]);
        }
    }
 }
 /// Test L1 loss computation
 #[cfg(all(test, feature = "disabled_tests"))]
 mod loss_tests {
    use super::*;
    #[test]
    fn test_l1_loss_computation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 2;
        let num_patches = 10;
        let patch_dim = 768;
        let predicted = Tensor::randn(vec![batch_size, num_patches, patch_dim], DType::F32, &device).unwrap();
        let target = Tensor::randn(vec![batch_size, num_patches, patch_dim], DType::F32, &device).unwrap();
        let loss = SimMIMLoss::compute_l1_loss(&predicted, &target).unwrap();
        // Loss should be a scalar
        assert_eq!(loss.shape(), vec![]);
        // Loss should be non-negative
        let loss_value = loss.to_vec::<f32>().unwrap()[0];
        assert!(loss_value >= 0.0);
    }
    #[test]
    fn test_l1_loss_identical_inputs() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let input = Tensor::ones(vec![2, 5, 100], DType::F32, &device).unwrap();
        let loss = SimMIMLoss::compute_l1_loss(&input, &input).unwrap();
        let loss_value = loss.to_vec::<f32>().unwrap()[0];
        // Loss should be zero for identical inputs
        assert!(loss_value < 1e-6);
    }
    #[test]
    fn test_l1_loss_known_values() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        // Create tensors with known values for deterministic testing
        let predicted = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![1, 3, 1], DType::F32, &device).unwrap();
        let target = Tensor::from_data(vec![2.0, 3.0, 1.0], vec![1, 3, 1], DType::F32, &device).unwrap();
        let loss = SimMIMLoss::compute_l1_loss(&predicted, &target).unwrap();
        let loss_value = loss.to_vec::<f32>().unwrap()[0];
        // Expected L1 loss: mean(|1-2| + |2-3| + |3-1|) = mean(1 + 1 + 2) = 4/3
        let expected = 4.0 / 3.0;
        assert!((loss_value - expected).abs() < 1e-6);
    }
    #[test]
    fn test_l1_loss_with_normalization() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let batch_size = 2;
        let num_patches = 8;
        let patch_dim = 768;
        let predicted = Tensor::randn(vec![batch_size, num_patches, patch_dim], DType::F32, &device).unwrap();
        let target = Tensor::randn(vec![batch_size, num_patches, patch_dim], DType::F32, &device).unwrap();
        let loss_normalized = SimMIMLoss::compute_l1_loss_normalized(&predicted, &target).unwrap();
        let loss_regular = SimMIMLoss::compute_l1_loss(&predicted, &target).unwrap();
        // Both should be scalars
        assert_eq!(loss_normalized.shape(), vec![]);
        assert_eq!(loss_regular.shape(), vec![]);
        // Normalized loss might be different depending on implementation
        let norm_value = loss_normalized.to_vec::<f32>().unwrap()[0];
        let reg_value = loss_regular.to_vec::<f32>().unwrap()[0];
        assert!(norm_value >= 0.0);
        assert!(reg_value >= 0.0);
    }
 }
 /// Test main SimMIM trainer
 #[cfg(all(test, feature = "disabled_tests"))]
 mod trainer_tests {
    use super::*;
    #[test]
    fn test_simmim_trainer_creation() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        assert_eq!(trainer.device(), &device);
        assert!(trainer.is_training());
    }
    #[test]
    fn test_simmim_config_defaults() {
        let config = SimMIMConfig::default();
        assert_eq!(config.mask_ratio, 0.6);
        assert_eq!(config.mask_patch_size, 32);
        assert_eq!(config.prediction_head, PredictionHeadType::Linear);
        assert_eq!(config.loss_type, LossType::L1);
        assert!(!config.norm_pix_loss);
        assert_eq!(config.encoder_dim, 768);
        assert_eq!(config.patch_size, 16);
    }
    #[test]
    fn test_simmim_config_with_methods() {
        let config = SimMIMConfig::default()
            .with_mask_ratio(0.75)
            .with_encoder_type(EncoderType::Swin)
            .with_norm_pix_loss(true);
        assert_eq!(config.mask_ratio, 0.75);
        assert_eq!(config.encoder_type, EncoderType::Swin);
        assert!(config.norm_pix_loss);
    }
    #[test]
    fn test_simmim_training_step() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let mut trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        let batch_size = 2;
        let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device).unwrap();
        let result = trainer.train_step(&images, Some(42)).unwrap();
        // Check training result structure
        assert_eq!(result.loss.shape(), vec![]);
        assert!(result.num_masked_patches > 0);
        assert!(result.mask_ratio > 0.0 && result.mask_ratio <= 1.0);
        let loss_value = result.loss.to_vec::<f32>().unwrap()[0];
        assert!(loss_value >= 0.0);
    }
    #[test]
    fn test_simmim_eval_mode() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let mut trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        // Switch to eval mode
        trainer.eval();
        assert!(!trainer.is_training());
        // Extract features should work in eval mode
        let images = Tensor::randn(vec![1, 3, 224, 224], DType::F32, &device).unwrap();
        let features = trainer.extract_features(&images).unwrap();
        // Should return feature tensor
        assert_eq!(features.shape()[0], 1); // batch size
        assert!(features.shape().len() >= 2); // at least [batch, features]
    }
    #[test]
    fn test_simmim_different_encoders() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let encoder_types = vec![EncoderType::ViT, EncoderType::Swin];
        for encoder_type in encoder_types {
            let config = SimMIMConfig::default()
                .with_encoder_type(encoder_type.clone());
            let trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
            assert_eq!(trainer.config().encoder_type, encoder_type);
        }
    }
    #[test]
    fn test_simmim_swin_specific_config() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        // Swin transformer should use different default mask ratio
        let config = SimMIMConfig::for_swin();
        assert_eq!(config.encoder_type, EncoderType::Swin);
        assert_eq!(config.mask_ratio, 0.32); // Lower mask ratio for hierarchical architecture
        assert_eq!(config.mask_patch_size, 32); // Larger patch size for hierarchical masking
    }
    #[test]
    fn test_simmim_vit_specific_config() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::for_vit();
        assert_eq!(config.encoder_type, EncoderType::ViT);
        assert_eq!(config.mask_ratio, 0.6); // Standard mask ratio for ViT
        assert_eq!(config.mask_patch_size, 16); // Standard patch size
    }
    #[test] 
    fn test_simmim_reproducibility() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let mut trainer1 = SimMIMTrainer::new(config.clone(), 3, 224, &device).unwrap();
        let mut trainer2 = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        let images = Tensor::randn(vec![2, 3, 224, 224], DType::F32, &device).unwrap();
        // Same seed should produce same results
        let result1 = trainer1.train_step(&images, Some(42)).unwrap();
        let result2 = trainer2.train_step(&images, Some(42)).unwrap();
        let loss1 = result1.loss.to_vec::<f32>().unwrap()[0];
        let loss2 = result2.loss.to_vec::<f32>().unwrap()[0];
        // Should be very close (within floating point precision)
        assert!((loss1 - loss2).abs() < 1e-5);
        assert_eq!(result1.num_masked_patches, result2.num_masked_patches);
    }
 }
 /// Test SimMIM integration with other components
 #[cfg(all(test, feature = "disabled_tests"))]
 mod integration_tests {
    use super::*;
    #[test]
    fn test_simmim_with_different_image_sizes() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let image_sizes = vec![224, 256, 384];
        for size in image_sizes {
            let trainer = SimMIMTrainer::new(config.clone(), 3, size, &device);
            assert!(trainer.is_ok(), "Failed to create trainer for image size {}", size);
            let trainer = trainer.unwrap();
            let images = Tensor::randn(vec![1, 3, size, size], DType::F32, &device).unwrap();
            let features = trainer.extract_features(&images);
            assert!(features.is_ok(), "Failed to extract features for image size {}", size);
        }
    }
    #[test] 
    fn test_simmim_memory_efficiency() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let mut trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        // Test with larger batch to ensure memory efficiency
        let batch_size = 8;
        let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device).unwrap();
        let result = trainer.train_step(&images, Some(42));
        assert!(result.is_ok(), "Failed to handle larger batch size");
    }
    #[test]
    fn test_simmim_patch_reconstruction_consistency() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let config = SimMIMConfig::default();
        let trainer = SimMIMTrainer::new(config, 3, 224, &device).unwrap();
        // Test that patch dimensions are consistent
        let images = Tensor::randn(vec![2, 3, 224, 224], DType::F32, &device).unwrap();
        // Extract patch information for validation
        let image_size = 224;
        let patch_size = trainer.config().patch_size;
        let num_patches = (image_size / patch_size).pow(2);
        assert!(num_patches > 0);
        assert_eq!(num_patches, 196); // 14x14 patches for 224x224 image with 16x16 patches
    }
 }
 /// Test error handling and edge cases
 #[cfg(all(test, feature = "disabled_tests"))]
 mod error_tests {
    use super::*;
    #[test]
    fn test_invalid_mask_ratio() {
        // Test mask ratios outside valid range
        let invalid_ratios = vec![-0.1, 1.1, 2.0];
        for ratio in invalid_ratios {
            let masker = RandomPatchMasker::new(ratio);
            // Should clamp to valid range [0.0, 1.0]
            assert!(masker.mask_ratio() >= 0.0 && masker.mask_ratio() <= 1.0);
        }
    }
    #[test]
    fn test_empty_tensor_handling() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let masker = RandomPatchMasker::new(0.5);
        // Test with empty tensor - should handle gracefully
        let empty_patches = Tensor::zeros(vec![0, 0, 768], DType::F32, &device).unwrap();
        let result = masker.mask_patches(&empty_patches, Some(42));
        // Should either handle gracefully or return appropriate error
        assert!(result.is_err() || result.unwrap().mask_indices.is_empty());
    }
    #[test]
    fn test_mismatched_tensor_dimensions() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let predicted = Tensor::randn(vec![2, 10, 768], DType::F32, &device).unwrap();
        let target = Tensor::randn(vec![2, 12, 768], DType::F32, &device).unwrap(); // Different sequence length
        let loss = SimMIMLoss::compute_l1_loss(&predicted, &target);
        assert!(loss.is_err(), "Should fail with mismatched dimensions");
    }
    #[test]
    fn test_zero_patches() {
        let device = Device::cuda(0).unwrap_or(Device::default());
        let masker = RandomPatchMasker::new(0.5);
        let patches = Tensor::randn(vec![1, 0, 768], DType::F32, &device).unwrap();
        let result = masker.mask_patches(&patches, Some(42));
        // Should handle zero patches gracefully
        assert!(result.is_err() || result.unwrap().visible_indices[0].is_empty());
    }
 }
@@ -0,0 +1,37 @@
 # Duplicate-implementation consolidation notes
 Recorded during the 2026-07-09 wiring/dead-code sweep. These are known
 duplications that were **not** consolidated in that pass because doing so
 requires rewriting call sites; new code should target the canonical
 implementation listed here.
 ## Mixture of Experts
 Canonical: `rtx-transformers/src/layers/mixture_of_experts/`
 Competing implementations still present (divergent internal APIs):
 - `rtx-transformers/src/layers/metal_moe.rs` (Metal-specific)
 - `rtx-transformers/src/modular/router.rs` (routing only)
 - `rtx-transformers/src/architectures/glam.rs` (architecture-embedded)
 Deleted in the sweep (were orphaned, never declared by `mod`):
 `layers/moe_layer.rs`, `layers/moe_integration.rs`.
 ## Flash Attention
 Canonical: the `rtx-flash-attention` crate (v2+v3 kernels, CUDA + Metal).
 Reimplementations inside rtx-transformers that should eventually delegate:
 - `src/tensor_core_kernels.rs`
 - `src/revolutionary/orchestrator_core.rs`
 - `src/ssl/jepa_gpu.rs`
 - `src/training/training_loop.rs`
 ## Speculative decoding
 Layering is intentional (documented in `rtx-inference/src/speculative/mod.rs`):
 `speculative/` is the orchestration layer (traits/configs/trees);
 `rtx-inference/src/medusa.rs` and `src/lookahead.rs` are the concrete
 implementations. Not a duplication to remove, but keep the two `MedusaConfig`
 types (orchestration vs implementation, re-exported as `MedusaHeadsConfig`)
 from drifting.