fix multiple bugs

app.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import ast
 import asyncio
+import io
 import os
 import queue
+import sys
 import traceback
 from pathlib import Path
 from typing import Any, Dict
@@ -205,6 +207,22 @@ def _make_exec_env(tracer: Tracer) -> Dict[str, Any]:
     }
 
 
+class PrintCapture(io.StringIO):
+    """Captures print output and sends it to the tracer."""
+    def __init__(self, tracer: Tracer):
+        super().__init__()
+        self.tracer = tracer
+        
+    def write(self, text: str) -> int:
+        # Send non-empty text to tracer
+        if text and text.strip():
+            self.tracer.print(text.rstrip('\n'))
+        return len(text)
+    
+    def flush(self):
+        pass
+
+
 def _run_user_code(code: str, tracer: Tracer) -> None:
     # 1. Transform code to auto-capture variable names
     transformed_code = transform_code(code)
@@ -212,14 +230,19 @@ def _run_user_code(code: str, tracer: Tracer) -> None:
     # 2. Setup Environment
     env = _make_exec_env(tracer)
 
-    # 3. Instrument Tensor/Layer classes to talk to our tracer
+    # 3. Redirect stdout to capture print statements
+    old_stdout = sys.stdout
+    sys.stdout = PrintCapture(tracer)
+
+    # 4. Instrument Tensor/Layer classes to talk to our tracer
     with Instrumentor(tracer):
         try:
-            # 4. Execute transformed code
+            # 5. Execute transformed code
             exec(transformed_code, env, env)
         except Exception:
             tracer.error(traceback.format_exc())
         finally:
+            sys.stdout = old_stdout
             tracer.done()
 
 

instrumentation.py

@@ -102,48 +102,62 @@ class Instrumentor:
         setattr(Tensor, method_name, wrapped)
 
     def _wrap_layer_forward(self):
-        """Wraps Layer.__call__ to emit trace events for layer operations."""
-        if not hasattr(Layer, '__call__'):
-            return
-            
-        original = Layer.__call__
-        self.original_methods["Layer.__call__"] = original
+        """Wraps Layer.__call__ and concrete layer forward methods."""
         instrumentor = self  # Capture reference for closure
-
-        def wrapped(instance, x, *args, **kwargs):
-            # Set flag to suppress internal tensor op tracing
-            was_inside = instrumentor._inside_layer
-            instrumentor._inside_layer = True
-            try:
-                # Execute original logic
-                result = original(instance, x, *args, **kwargs)
-            finally:
-                instrumentor._inside_layer = was_inside
-
-            # Get the layer name
-            layer_name = instance.__class__.__name__
-
-            # Build inputs list - for Linear, include weight and bias
-            inputs = [x]
-            meta = {'layer_type': layer_name}
-            
-            # For Linear layers, include weight and bias for visualization
-            if layer_name == 'Linear':
-                if hasattr(instance, 'weight'):
-                    #inputs.append(instance.weight)
-                    inputs.insert(0, instance.weight)
-
-                    meta['has_weight'] = True
-                if hasattr(instance, 'bias') and instance.bias is not None:
-                    inputs.append(instance.bias)
-                    meta['has_bias'] = True
-
-            # Emit op event for the layer
-            instrumentor.tracer.op(layer_name.lower(), inputs, result, meta)
-
-            return result
-
-        setattr(Layer, "__call__", wrapped)
+        
+        # Import Linear here to wrap its forward method specifically
+        from tinytorch.core.layers import Linear, Dropout
+
+        def make_layer_wrapper(layer_cls):
+            def make_wrapped(orig):
+                def wrapped(instance, x, *args, **kwargs):
+                    # Set flag to suppress internal tensor op tracing
+                    was_inside = instrumentor._inside_layer
+                    instrumentor._inside_layer = True
+                    try:
+                        # Execute original logic
+                        result = orig(instance, x, *args, **kwargs)
+                    finally:
+                        instrumentor._inside_layer = was_inside
+
+                    # Only emit event if this is the outermost call (prevent double-tracing)
+                    if not was_inside:
+                        # Get the layer name
+                        layer_name = instance.__class__.__name__
+
+                        # Build inputs list - for Linear, include weight and bias
+                        inputs = [x]
+                        meta = {'layer_type': layer_name}
+                        
+                        # For Linear layers, include weight and bias for visualization
+                        if layer_name == 'Linear':
+                            if hasattr(instance, 'weight'):
+                                inputs.insert(0, instance.weight)
+                                meta['has_weight'] = True
+                            if hasattr(instance, 'bias') and instance.bias is not None:
+                                inputs.append(instance.bias)
+                                meta['has_bias'] = True
+
+                        # Emit op event for the layer
+                        instrumentor.tracer.op(layer_name.lower(), inputs, result, meta)
+
+                    return result
+                return wrapped
+            return make_wrapped
+
+        # Wrap Layer.__call__
+        if hasattr(Layer, '__call__'):
+            original_call = Layer.__call__
+            self.original_methods["Layer.__call__"] = original_call
+            setattr(Layer, "__call__", make_layer_wrapper(Layer)(original_call))
+        
+        # Wrap concrete layer forward methods (Linear, Dropout, etc.)
+        for layer_cls in [Linear, Dropout]:
+            if hasattr(layer_cls, 'forward'):
+                cls_name = layer_cls.__name__
+                original_forward = layer_cls.forward
+                self.original_methods[f"{cls_name}.forward"] = original_forward
+                setattr(layer_cls, "forward", make_layer_wrapper(layer_cls)(original_forward))
 
     def _wrap_activation(self, activation_cls):
         """Wraps an activation class's __call__ and forward methods."""

static/index.html

@@ -107,6 +107,90 @@
     #run:hover { background: #2563eb; }
     #run:disabled { background: #475569; cursor: not-allowed; }
     #error { color: #ef4444; font-size: 12px; min-height: 20px; font-family: monospace; white-space: pre-wrap; }
+    
+    /* Console output area */
+    #console-container {
+      margin-top: 12px;
+      border: 1px solid #334155;
+      border-radius: 6px;
+      background: #0a0f1a;
+      overflow: hidden;
+      display: flex;
+      flex-direction: column;
+      min-height: 60px;
+      height: 150px;
+      flex-shrink: 0;
+    }
+    #console-resizer {
+      height: 6px;
+      background: transparent;
+      cursor: ns-resize;
+      position: relative;
+      flex-shrink: 0;
+    }
+    #console-resizer::after {
+      content: '';
+      position: absolute;
+      left: 50%;
+      top: 50%;
+      transform: translate(-50%, -50%);
+      width: 40px;
+      height: 3px;
+      background: #475569;
+      border-radius: 2px;
+      transition: background 0.2s;
+    }
+    #console-resizer:hover::after {
+      background: #64748b;
+    }
+    #console-header {
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+      padding: 6px 10px;
+      background: #1e293b;
+      border-bottom: 1px solid #334155;
+      font-size: 11px;
+      font-weight: 600;
+      color: #94a3b8;
+      letter-spacing: 0.5px;
+    }
+    #console-clear {
+      background: none;
+      border: none;
+      color: #64748b;
+      cursor: pointer;
+      font-size: 12px;
+      padding: 2px 6px;
+      border-radius: 3px;
+    }
+    #console-clear:hover {
+      background: #334155;
+      color: #f1f5f9;
+    }
+    #console-output {
+      flex: 1;
+      overflow-y: auto;
+      padding: 8px 10px;
+      font-family: 'Fira Code', 'JetBrains Mono', 'Consolas', monospace;
+      font-size: 12px;
+      line-height: 1.5;
+      color: #e2e8f0;
+      white-space: pre-wrap;
+      word-break: break-word;
+    }
+    #console-output .console-line {
+      margin: 2px 0;
+    }
+    #console-output .console-line.error {
+      color: #f87171;
+    }
+    #console-output .console-line.success {
+      color: #4ade80;
+    }
+    #console-output .console-line.info {
+      color: #60a5fa;
+    }
 
     #pdfSidebar {
       flex: 1 1 340px;
@@ -533,6 +617,32 @@
       color: #ff4444;
       text-shadow: 0 0 4px rgba(255, 68, 68, 0.5);
     }
+    
+    /* Console in retro mode */
+    body.retro-mode #console-resizer::after {
+      background: rgba(0, 255, 65, 0.3);
+    }
+    body.retro-mode #console-resizer:hover::after {
+      background: rgba(0, 255, 65, 0.6);
+    }
+    body.retro-mode #console-container {
+      border-color: rgba(0, 255, 65, 0.3);
+      background: rgba(0, 10, 5, 0.8);
+    }
+    body.retro-mode #console-header {
+      background: rgba(0, 40, 20, 0.6);
+      border-bottom-color: rgba(0, 255, 65, 0.2);
+      color: #00ff41;
+      text-shadow: 0 0 4px rgba(0, 255, 65, 0.5);
+    }
+    body.retro-mode #console-output {
+      color: #00ff41;
+      text-shadow: 0 0 2px rgba(0, 255, 65, 0.3);
+    }
+    body.retro-mode #console-output .console-line.error {
+      color: #ff4444;
+      text-shadow: 0 0 4px rgba(255, 68, 68, 0.5);
+    }
 
     /* Status dot glow */
     body.retro-mode .status-dot {
@@ -612,6 +722,14 @@ box("Layer 3: Linear + Softmax", [a2, z3, y_pred], "4")
 box("Loss Computation", [y_pred, target_probs, loss], "6")
     </textarea>
     <button id="run">Run</button>
+    <div id="console-resizer"></div>
+    <div id="console-container">
+      <div id="console-header">
+        <span>CONSOLE OUTPUT</span>
+        <button id="console-clear" title="Clear console">✕</button>
+      </div>
+      <div id="console-output"></div>
+    </div>
     <div id="error"></div>
   </div>
   <div class="v-resizer" data-left="editor"></div>
@@ -724,6 +842,7 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
         nextGroupId = 1;
         nextBoxId = 1;
         document.getElementById('error').textContent = '';
+        clearConsole();
         break;
       case 'tensor':
         tensors[msg.id] = { id: msg.id, shape: msg.shape, data: msg.data, name: msg.name };
@@ -745,12 +864,30 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
         break;
       case 'error':
         document.getElementById('error').textContent = msg.message;
+        appendToConsole(msg.message, 'error');
+        break;
+      case 'print':
+        appendToConsole(msg.text, msg.type || 'info');
         break;
       case 'done':
         layoutAndRender();
         break;
     }
   }
+  
+  // ==================== Console Output ====================
+  function appendToConsole(text, type = 'info') {
+    const consoleOutput = document.getElementById('console-output');
+    const line = document.createElement('div');
+    line.className = 'console-line ' + type;
+    line.textContent = text;
+    consoleOutput.appendChild(line);
+    consoleOutput.scrollTop = consoleOutput.scrollHeight;
+  }
+  
+  function clearConsole() {
+    document.getElementById('console-output').innerHTML = '';
+  }
 
   // ==================== Layout ====================
   function layoutAndRender() {
@@ -1495,6 +1632,44 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
     });
   }
 
+  // ==================== Transpose Hover Highlighting ====================
+  // For transpose, output[i,j] corresponds to input[j,i]
+  function setupTransposeHover(inputCard, outputCard) {
+    const outputTable = outputCard.querySelector('.matrix-table');
+    const inputTable = inputCard.querySelector('.matrix-table');
+    
+    if (!outputTable || !inputTable) return;
+    
+    const outputCells = outputTable.querySelectorAll('td[data-row][data-col]');
+    
+    outputCells.forEach(cell => {
+      cell.style.cursor = 'pointer';
+      
+      cell.addEventListener('mouseenter', () => {
+        const row = cell.dataset.row;
+        const col = cell.dataset.col;
+        
+        // Highlight the output cell
+        cell.classList.add('highlight-cell');
+        
+        // For transpose: input[col, row] -> output[row, col]
+        // So to find the source, swap row and col
+        const inputCell = inputTable.querySelector(`td[data-row="${col}"][data-col="${row}"]`);
+        if (inputCell) {
+          inputCell.classList.add('highlight-cell');
+        }
+      });
+      
+      cell.addEventListener('mouseleave', () => {
+        // Remove all highlights
+        cell.classList.remove('highlight-cell');
+        inputTable.querySelectorAll('td.highlight-cell').forEach(td => {
+          td.classList.remove('highlight-cell');
+        });
+      });
+    });
+  }
+
   // ==================== Matmul Hover Highlighting ====================
   function setupMatmulHover(wrapper, leftCard, topCard, resultCard) {
     const resultTable = resultCard.querySelector('.matrix-table');
@@ -1703,7 +1878,7 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
     
     if (isMatmul) {
       // Matrix multiplication: keep grid layout for proper alignment
-      const grid = document.createElement('div');
+    const grid = document.createElement('div');
       grid.className = 'layout-grid layout-binary';
     
       const left = createMatrixCard(inputs[0].name || inputs[0].id, inputs[0], 'auto', 'left');
@@ -1748,7 +1923,7 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
     }
     
     const inputCards = [];
-    if (inputs[0]) {
+      if (inputs[0]) {
       const inputLabel = getInputDisplayName(inputs[0]);
       const inputCard = createMatrixCard(inputLabel, inputs[0], 'auto', 'input');
       grid.appendChild(inputCard);
@@ -1757,10 +1932,16 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
     const outputCard = createMatrixCard(output.name || type, output, outputOrientation, 'output');
     grid.appendChild(outputCard);
     
-    // Add element-wise hover highlighting for activations and element-wise ops
-    const isElementwiseUnary = ['relu', 'sigmoid', 'tanh', 'gelu', 'softmax', 'logsoftmax'].includes(type);
-    if (isElementwiseUnary && inputCards.length > 0) {
-      setupElementwiseHover(inputCards, outputCard);
+    // Add hover highlighting based on operation type
+    if (type === 'transpose' && inputCards.length > 0) {
+      // Transpose: output[i,j] corresponds to input[j,i]
+      setupTransposeHover(inputCards[0], outputCard);
+    } else {
+      // Element-wise activations: output[i,j] corresponds to input[i,j]
+      const isElementwiseUnary = ['relu', 'sigmoid', 'tanh', 'gelu', 'softmax', 'logsoftmax'].includes(type);
+      if (isElementwiseUnary && inputCards.length > 0) {
+        setupElementwiseHover(inputCards, outputCard);
+      }
     }
     
     container.appendChild(grid);
@@ -2172,6 +2353,39 @@ box("Loss Computation", [y_pred, target_probs, loss], "6")
   document.getElementById('code').onkeydown = e => {
     if ((e.ctrlKey || e.metaKey) && e.key === 'Enter') { e.preventDefault(); runCode(); }
   };
+  document.getElementById('console-clear').onclick = clearConsole;
+
+  // ==================== Console Resizer ====================
+  (function() {
+    const resizer = document.getElementById('console-resizer');
+    const consoleContainer = document.getElementById('console-container');
+    let startY, startHeight;
+    
+    resizer.addEventListener('mousedown', (e) => {
+      e.preventDefault();
+      startY = e.clientY;
+      startHeight = consoleContainer.offsetHeight;
+      
+      document.addEventListener('mousemove', onMouseMove);
+      document.addEventListener('mouseup', onMouseUp);
+      document.body.style.cursor = 'ns-resize';
+      document.body.style.userSelect = 'none';
+    });
+    
+    function onMouseMove(e) {
+      // Dragging up increases height (startY - e.clientY is positive when moving up)
+      const deltaY = startY - e.clientY;
+      const newHeight = Math.max(60, Math.min(400, startHeight + deltaY));
+      consoleContainer.style.height = newHeight + 'px';
+    }
+    
+    function onMouseUp() {
+      document.removeEventListener('mousemove', onMouseMove);
+      document.removeEventListener('mouseup', onMouseUp);
+      document.body.style.cursor = '';
+      document.body.style.userSelect = '';
+    }
+  })();
 
   // ==================== Init ====================
   // Restore retro mode preference

tracer.py

@@ -176,5 +176,9 @@ class Tracer:
     def error(self, message: str) -> None:
         self.sink.emit(TraceEvent("error", {"message": str(message)}).asdict())
 
+    def print(self, text: str, msg_type: str = "info") -> None:
+        """Emit a print event for console output."""
+        self.sink.emit(TraceEvent("print", {"text": str(text), "type": msg_type}).asdict())
+
     def done(self) -> None:
         self.sink.emit(TraceEvent("done", {}).asdict())