Upload 2 files

ai.py

@@ -47,6 +47,39 @@ class AIAccelerator:
             bits=32,
             bandwidth_tbps=drift_velocity / 1e-12  # Bandwidth scaled to electron drift speed
         )
+        self.tensor_cores_initialized = False
+        
+    def initialize_tensor_cores(self):
+        """Initialize tensor cores and verify they're ready for computation"""
+        if self.tensor_cores_initialized:
+            return True
+            
+        try:
+            # Verify tensor core array is properly initialized
+            if not hasattr(self, 'tensor_core_array') or self.tensor_core_array is None:
+                raise RuntimeError("Tensor core array not properly initialized")
+                
+            # Initialize tensor cores if needed
+            if hasattr(self.tensor_core_array, 'initialize'):
+                self.tensor_core_array.initialize()
+                
+            # Verify VRAM access
+            if self.vram is None:
+                raise RuntimeError("VRAM not properly configured")
+                
+            # Test tensor core functionality with a small computation
+            test_input = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+            test_result = self.tensor_core_array.matmul(test_input, test_input)
+            if test_result is None or test_result.size == 0:
+                raise RuntimeError("Tensor core test computation failed")
+                
+            self.tensor_cores_initialized = True
+            return True
+            
+        except Exception as e:
+            print(f"Failed to initialize tensor cores: {str(e)}")
+            self.tensor_cores_initialized = False
+            return False
         
         # AI operation statistics
         self.operations_performed = 0

virtual_vram.py

@@ -23,9 +23,37 @@ class VirtualVRAM:
         }
         self.store_vram_state()
         
-    def store_vram_state(self):
-        """Store VRAM state in WebSocket storage"""
-        self.storage.store_state("vram", "state", self.vram_state)
+    def store_vram_state(self, max_retries=3):
+        """Store VRAM state in WebSocket storage with retry logic"""
+        for attempt in range(max_retries):
+            try:
+                # Wait for connection if needed
+                if not self.storage.wait_for_connection(timeout=5):
+                    print(f"Waiting for WebSocket connection (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(1)
+                    continue
+                    
+                # Ensure state is JSON serializable
+                safe_state = {
+                    "total_size": str(self.vram_state["total_size"]) if isinstance(self.vram_state["total_size"], float) and self.vram_state["total_size"] == float('inf') else self.vram_state["total_size"],
+                    "allocated": self.vram_state["allocated"],
+                    "blocks": self.vram_state["blocks"],
+                    "memory_map": self.vram_state["memory_map"],
+                    "is_unlimited": self.vram_state["is_unlimited"]
+                }
+                
+                success = self.storage.store_state("vram", "state", safe_state)
+                if success:
+                    return True
+                    
+                print(f"Failed to store VRAM state (attempt {attempt + 1}/{max_retries})")
+                time.sleep(1)
+                
+            except Exception as e:
+                print(f"Error storing VRAM state (attempt {attempt + 1}/{max_retries}): {str(e)}")
+                time.sleep(1)
+                
+        raise RuntimeError("Failed to store VRAM state after multiple attempts")
         
     def allocate_block(self, size: int, block_id: Optional[str] = None) -> str:
         """Allocate a block of VRAM"""