Upload 7 files

README.md

@@ -1,9 +1,56 @@
 ---
 title: Neural Network Playground
-emoji: 🌖
+emoji: 🧠
 colorFrom: pink
 colorTo: blue
 sdk: static
-pinned: false
+pinned: true
+license: mit
 ---
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# Neural Network Playground
+
+![Neural Network Playground](https://raw.githubusercontent.com/huggingface/hub-docs/main/static/logos/huggingface_logo-noborder.svg)
+
+## Introduction
+
+Neural Network Playground is an interactive visualization tool that helps you understand how neural networks work. Built with plain HTML, CSS, and JavaScript, it allows you to:
+
+- Create custom neural network architectures by dragging and dropping different types of layers
+- Connect layers and see how data flows through the network
+- View input and output shapes for each layer
+- Visualize layer parameters and configurations
+
+## Features
+
+- **Interactive Interface**: Drag and drop nodes to create neural networks
+- **Shape Information**: See input and output shapes for each node
+- **Detailed Parameters**: View kernel size, stride, and padding for applicable layers
+- **Layer Types**:
+  - Input Layer
+  - Hidden Layer
+  - Output Layer
+  - Convolutional Layer
+  - Pooling Layer
+  - Linear Regression Layer
+
+## How to Use
+
+1. Drag components from the left panel onto the canvas
+2. Connect them by dragging from output (right) ports to input (left) ports
+3. Double-click on nodes to edit their properties
+4. Use the network settings to adjust learning rate and activation functions
+
+## Technical Details
+
+The playground visualizes how neural networks process data and helps users understand concepts like:
+
+- Shape transformations between layers
+- Parameter calculations
+- The effects of different layer configurations
+
+This is an educational tool designed to make neural networks more accessible and understandable.
+
+## License
+
+MIT

js/drag-drop.js

@@ -16,6 +16,17 @@ function initializeDragAndDrop() {
         connections: []
     };
     
+    // Helper function to format numbers with K, M, B suffixes
+    function formatNumber(num) {
+        if (num === 0) return '0';
+        if (!num) return 'N/A';
+        
+        if (num >= 1e9) return (num / 1e9).toFixed(2) + 'B';
+        if (num >= 1e6) return (num / 1e6).toFixed(2) + 'M';
+        if (num >= 1e3) return (num / 1e3).toFixed(2) + 'K';
+        return num.toString();
+    }
+    
     // Add event listeners to draggable items
     nodeItems.forEach(item => {
         item.addEventListener('dragstart', handleDragStart);
@@ -111,7 +122,7 @@ function initializeDragAndDrop() {
                     inputShape = 'Connect input';
                     outputShape = 'Depends on input';
                     // Create parameter string
-                    parameters = `In: ?, Out: ${nodeConfig.filters}\nKernel: ${nodeConfig.kernelSize.join('×')}\nStride: ${nodeConfig.strides.join('×')}\nPadding: ${nodeConfig.padding}`;
+                    parameters = `Kernel: ${nodeConfig.kernelSize.join('×')}\nStride: ${nodeConfig.strides.join('×')}\nPadding: ${nodeConfig.padding}`;
                     break;
                 case 'pool':
                     const poolCount = document.querySelectorAll('.canvas-node[data-type="pool"]').length;
@@ -127,11 +138,6 @@ function initializeDragAndDrop() {
                     parameters = 'N/A';
             }
             
-            // Create node header
-            const nodeHeader = document.createElement('div');
-            nodeHeader.className = 'node-header';
-            nodeHeader.textContent = nodeName;
-            
             // Create node content
             const nodeContent = document.createElement('div');
             nodeContent.className = 'node-content';
@@ -147,25 +153,71 @@ function initializeDragAndDrop() {
             // Add parameters section
             const paramsSection = document.createElement('div');
             paramsSection.className = 'params-section';
-            paramsSection.innerHTML = `<pre class="params-display">${parameters}</pre>`;
-            
-            // Add connection ports
-            const inputPort = document.createElement('div');
-            inputPort.className = 'port input-port';
-            inputPort.setAttribute('data-port-type', 'input');
-            
-            const outputPort = document.createElement('div');
-            outputPort.className = 'port output-port';
-            outputPort.setAttribute('data-port-type', 'output');
+            paramsSection.innerHTML = `
+                <div class="params-details">${parameters}</div>
+                <div class="node-parameters">Params: ${nodeConfig.parameters !== undefined ? formatNumber(nodeConfig.parameters) : '?'}</div>
+            `;
             
-            // Assemble the node
+            // Assemble content
             nodeContent.appendChild(shapeInfo);
             nodeContent.appendChild(paramsSection);
             
-            canvasNode.appendChild(nodeHeader);
+            // Add dimensions section to show shapes compactly
+            const dimensionsSection = document.createElement('div');
+            dimensionsSection.className = 'node-dimensions';
+            
+            // Set dimensions text based on node type
+            let dimensionsText = '';
+            switch(nodeType) {
+                case 'input':
+                    dimensionsText = nodeConfig.shape.join(' × ');
+                    break;
+                case 'hidden':
+                case 'output':
+                    dimensionsText = nodeConfig.units.toString();
+                    break;
+                case 'conv':
+                    if (nodeConfig.inputShape && nodeConfig.outputShape) {
+                        dimensionsText = `${nodeConfig.inputShape.join('×')} → ${nodeConfig.outputShape.join('×')}`;
+                    } else {
+                        dimensionsText = `? → ${nodeConfig.filters} filters`;
+                    }
+                    break;
+                case 'pool':
+                    if (nodeConfig.inputShape && nodeConfig.outputShape) {
+                        dimensionsText = `${nodeConfig.inputShape.join('×')} → ${nodeConfig.outputShape.join('×')}`;
+                    } else {
+                        dimensionsText = `? → ?`;
+                    }
+                    break;
+                case 'linear':
+                    dimensionsText = `${nodeConfig.inputFeatures} → ${nodeConfig.outputFeatures}`;
+                    break;
+            }
+            dimensionsSection.textContent = dimensionsText;
+            
+            // Add node title for clearer identification
+            const nodeTitle = document.createElement('div');
+            nodeTitle.className = 'node-title';
+            nodeTitle.textContent = nodeName;
+            
+            // Add connection ports
+            const portIn = document.createElement('div');
+            portIn.className = 'node-port port-in';
+            
+            const portOut = document.createElement('div');
+            portOut.className = 'node-port port-out';
+            
+            // Assemble the node with the new structure
+            canvasNode.appendChild(nodeTitle);
+            canvasNode.appendChild(dimensionsSection);
             canvasNode.appendChild(nodeContent);
-            canvasNode.appendChild(inputPort);
-            canvasNode.appendChild(outputPort);
+            canvasNode.appendChild(portIn);
+            canvasNode.appendChild(portOut);
+            
+            // Store node data attributes for easier access
+            canvasNode.setAttribute('data-name', nodeName);
+            canvasNode.setAttribute('data-dimensions', dimensionsText);
             
             // Add node to the canvas
             canvas.appendChild(canvasNode);
@@ -175,10 +227,13 @@ function initializeDragAndDrop() {
             
             // Add event listeners for node manipulation
             canvasNode.addEventListener('mousedown', startDrag);
-            inputPort.addEventListener('mousedown', (e) => {
+            
+            // Update port event listeners for the new class names
+            portIn.addEventListener('mousedown', (e) => {
                 e.stopPropagation();
             });
-            outputPort.addEventListener('mousedown', (e) => {
+            
+            portOut.addEventListener('mousedown', (e) => {
                 e.stopPropagation();
                 startConnection(canvasNode, e);
             });
@@ -198,7 +253,11 @@ function initializeDragAndDrop() {
             networkLayers.layers.push({
                 id: layerId,
                 type: nodeType,
-                config: nodeConfig
+                name: nodeName,
+                position: { x, y },
+                dimensions: dimensionsText,
+                config: nodeConfig,
+                parameters: nodeConfig.parameters || 0
             });
             
             // Notify about network changes
@@ -251,12 +310,19 @@ function initializeDragAndDrop() {
         let x = e.clientX - canvasRect.left - offsetX;
         let y = e.clientY - canvasRect.top - offsetY;
         
-        // Constrain to canvas
-        x = Math.max(0, Math.min(canvasRect.width - draggedNode.offsetWidth, x));
-        y = Math.max(0, Math.min(canvasRect.height - draggedNode.offsetHeight, y));
+        // Constrain to canvas with better boundary checks
+        const nodeWidth = draggedNode.offsetWidth || 150; // Default width if not set
+        const nodeHeight = draggedNode.offsetHeight || 100; // Default height if not set
+        
+        // Ensure the node stays completely within the canvas
+        x = Math.max(0, Math.min(canvasRect.width - nodeWidth, x));
+        y = Math.max(0, Math.min(canvasRect.height - nodeHeight, y));
         
+        // Apply position with fixed sizing to prevent layout expansion
+        draggedNode.style.position = 'absolute';
         draggedNode.style.left = `${x}px`;
         draggedNode.style.top = `${y}px`;
+        draggedNode.style.width = `${nodeWidth}px`; // Maintain fixed width
         
         // Update node position in network layers
         const nodeId = draggedNode.getAttribute('data-id');
@@ -509,365 +575,4 @@ function initializeDragAndDrop() {
             const endY = targetPortRect.top + (targetPortRect.height / 2) - canvasRect.top;
             
             // Create the connection
-            const pathId = `connection-${sourceId}-${targetId}`;
-            const connectionPath = document.createElementNS('http://www.w3.org/2000/svg', 'path');
-            connectionPath.setAttribute('id', pathId);
-            connectionPath.setAttribute('class', 'connection-line');
-            
-            // Curved path (bezier)
-            const dx = Math.abs(endX - startX) * 0.7;
-            const path = `M ${startX} ${startY} C ${startX + dx} ${startY}, ${endX - dx} ${endY}, ${endX} ${endY}`;
-            connectionPath.setAttribute('d', path);
-            
-            // Add connection to SVG container
-            svgContainer.appendChild(connectionPath);
-            
-            // Add to connections
-            networkLayers.connections.push({
-                id: pathId,
-                source: sourceId,
-                target: targetId,
-                sourceType: sourceType,
-                targetType: targetType
-            });
-            
-            // Update input and output shapes
-            updateNodeShapes(sourceId, targetId);
-            
-            // Notify about connection
-            document.dispatchEvent(new CustomEvent('networkUpdated', {
-                detail: networkLayers
-            }));
-        }
-        
-        // Clean up
-        removePortHighlights();
-        if (connectionLine) {
-            connectionLine.remove();
-            connectionLine = null;
-        }
-        isConnecting = false;
-        startNode = null;
-    }
-    
-    // Update input and output shapes when connections are made
-    function updateNodeShapes(sourceId, targetId) {
-        const sourceNode = document.querySelector(`.canvas-node[data-id="${sourceId}"]`);
-        const targetNode = document.querySelector(`.canvas-node[data-id="${targetId}"]`);
-        
-        if (sourceNode && targetNode) {
-            const sourceConfig = sourceNode.layerConfig;
-            const targetConfig = targetNode.layerConfig;
-            
-            // Update the target's input shape based on the source's output shape
-            if (sourceConfig && targetConfig) {
-                // Calculate output shape based on node type
-                let outputShape;
-                switch (sourceNode.getAttribute('data-type')) {
-                    case 'input':
-                        outputShape = sourceConfig.shape;
-                        break;
-                    case 'hidden':
-                        outputShape = [sourceConfig.units];
-                        break;
-                    case 'output':
-                        outputShape = [sourceConfig.units];
-                        break;
-                    case 'conv':
-                        // For Conv2D, the output shape depends on the input and parameters
-                        // This is a simplified calculation
-                        if (targetConfig.inputShape) {
-                            const h = targetConfig.inputShape[0];
-                            const w = targetConfig.inputShape[1];
-                            const kh = sourceConfig.kernelSize[0];
-                            const kw = sourceConfig.kernelSize[1];
-                            const sh = sourceConfig.strides[0];
-                            const sw = sourceConfig.strides[1];
-                            const padding = sourceConfig.padding;
-                            
-                            let outHeight, outWidth;
-                            if (padding === 'same') {
-                                outHeight = Math.ceil(h / sh);
-                                outWidth = Math.ceil(w / sw);
-                            } else { // 'valid'
-                                outHeight = Math.ceil((h - kh + 1) / sh);
-                                outWidth = Math.ceil((w - kw + 1) / sw);
-                            }
-                            
-                            outputShape = [outHeight, outWidth, sourceConfig.filters];
-                        } else {
-                            outputShape = ['?', '?', sourceConfig.filters];
-                        }
-                        break;
-                    case 'pool':
-                        // For pooling, also depends on the input and parameters
-                        if (targetConfig.inputShape) {
-                            const h = targetConfig.inputShape[0];
-                            const w = targetConfig.inputShape[1];
-                            const c = targetConfig.inputShape[2];
-                            const ph = sourceConfig.poolSize[0];
-                            const pw = sourceConfig.poolSize[1];
-                            const sh = sourceConfig.strides[0];
-                            const sw = sourceConfig.strides[1];
-                            const padding = sourceConfig.padding;
-                            
-                            let outHeight, outWidth;
-                            if (padding === 'same') {
-                                outHeight = Math.ceil(h / sh);
-                                outWidth = Math.ceil(w / sw);
-                            } else { // 'valid'
-                                outHeight = Math.ceil((h - ph + 1) / sh);
-                                outWidth = Math.ceil((w - pw + 1) / sw);
-                            }
-                            
-                            outputShape = [outHeight, outWidth, c];
-                        } else {
-                            outputShape = ['?', '?', '?'];
-                        }
-                        break;
-                    case 'linear':
-                        outputShape = [sourceConfig.outputFeatures];
-                        break;
-                    default:
-                        outputShape = ['?', '?', '?'];
-                }
-                
-                // Update the target's input shape
-                targetConfig.inputShape = outputShape;
-                
-                // Update UI
-                updateNodeDisplayShapes(sourceNode, targetNode);
-            }
-        }
-    }
-    
-    // Update the displayed shapes in the UI
-    function updateNodeDisplayShapes(sourceNode, targetNode) {
-        if (sourceNode && targetNode) {
-            const sourceType = sourceNode.getAttribute('data-type');
-            const targetType = targetNode.getAttribute('data-type');
-            const sourceConfig = sourceNode.layerConfig;
-            const targetConfig = targetNode.layerConfig;
-            
-            // Update source node output shape display
-            const sourceOutputElem = sourceNode.querySelector('.output-shape');
-            if (sourceOutputElem && sourceConfig) {
-                let outputText;
-                switch (sourceType) {
-                    case 'input':
-                        outputText = `[${sourceConfig.shape.join(' × ')}]`;
-                        break;
-                    case 'hidden':
-                    case 'output':
-                        outputText = `[${sourceConfig.units}]`;
-                        break;
-                    case 'conv':
-                        if (sourceConfig.outputShape) {
-                            outputText = `[${sourceConfig.outputShape.join(' × ')}]`;
-                        } else {
-                            outputText = `[? × ? × ${sourceConfig.filters}]`;
-                        }
-                        break;
-                    case 'pool':
-                        if (sourceConfig.outputShape) {
-                            outputText = `[${sourceConfig.outputShape.join(' × ')}]`;
-                        } else {
-                            outputText = 'Depends on input';
-                        }
-                        break;
-                    case 'linear':
-                        outputText = `[${sourceConfig.outputFeatures}]`;
-                        break;
-                    default:
-                        outputText = 'Unknown';
-                }
-                sourceOutputElem.textContent = outputText;
-            }
-            
-            // Update target node input shape display
-            const targetInputElem = targetNode.querySelector('.input-shape');
-            if (targetInputElem && targetConfig && targetConfig.inputShape) {
-                targetInputElem.textContent = `[${targetConfig.inputShape.join(' × ')}]`;
-                
-                // Update parameters section
-                const targetParamsElem = targetNode.querySelector('.params-display');
-                if (targetParamsElem) {
-                    // Calculate and display parameters
-                    let paramsText = '';
-                    switch (targetType) {
-                        case 'hidden':
-                            const inputUnits = Array.isArray(targetConfig.inputShape) ? 
-                                targetConfig.inputShape.reduce((acc, val) => acc * val, 1) : 
-                                targetConfig.inputShape;
-                                
-                            const biasParams = targetConfig.useBias ? targetConfig.units : 0;
-                            const totalParams = (inputUnits * targetConfig.units) + biasParams;
-                            
-                            paramsText = `In: ${inputUnits}, Out: ${targetConfig.units}\nParams: ${totalParams.toLocaleString()}\nDropout: ${targetConfig.dropoutRate}`;
-                            break;
-                        case 'output':
-                            const outInputUnits = Array.isArray(targetConfig.inputShape) ? 
-                                targetConfig.inputShape.reduce((acc, val) => acc * val, 1) : 
-                                targetConfig.inputShape;
-                                
-                            const outBiasParams = targetConfig.useBias ? targetConfig.units : 0;
-                            const outTotalParams = (outInputUnits * targetConfig.units) + outBiasParams;
-                            
-                            paramsText = `In: ${outInputUnits}, Out: ${targetConfig.units}\nParams: ${outTotalParams.toLocaleString()}\nActivation: ${targetConfig.activation}`;
-                            break;
-                        case 'conv':
-                            const channels = targetConfig.inputShape[2] || '?';
-                            const kernelH = targetConfig.kernelSize[0];
-                            const kernelW = targetConfig.kernelSize[1];
-                            const kernelParams = kernelH * kernelW * channels * targetConfig.filters;
-                            const convBiasParams = targetConfig.useBias ? targetConfig.filters : 0;
-                            const convTotalParams = kernelParams + convBiasParams;
-                            
-                            paramsText = `In: ${channels}, Out: ${targetConfig.filters}\nKernel: ${targetConfig.kernelSize.join('×')}\nStride: ${targetConfig.strides.join('×')}\nPadding: ${targetConfig.padding}\nParams: ${convTotalParams.toLocaleString()}`;
-                            break;
-                        case 'pool':
-                            paramsText = `Pool size: ${targetConfig.poolSize.join('×')}\nStride: ${targetConfig.strides.join('×')}\nPadding: ${targetConfig.padding}\nParams: 0`;
-                            break;
-                        case 'linear':
-                            const linearInputs = targetConfig.inputFeatures;
-                            const linearBiasParams = targetConfig.useBias ? targetConfig.outputFeatures : 0;
-                            const linearTotalParams = (linearInputs * targetConfig.outputFeatures) + linearBiasParams;
-                            
-                            paramsText = `In: ${linearInputs}, Out: ${targetConfig.outputFeatures}\nParams: ${linearTotalParams.toLocaleString()}\nLearning Rate: ${targetConfig.learningRate}\nLoss: ${targetConfig.lossFunction}`;
-                            break;
-                    }
-                    
-                    targetParamsElem.textContent = paramsText;
-                }
-            }
-        }
-    }
-    
-    // Delete a node and its connections
-    function deleteNode(node) {
-        if (!node) return;
-        
-        const nodeId = node.getAttribute('data-id');
-        
-        // Remove all connections to/from this node
-        document.querySelectorAll(`.connection[data-source="${nodeId}"], .connection[data-target="${nodeId}"]`).forEach(conn => {
-            conn.parentNode.removeChild(conn);
-        });
-        
-        // Remove from network layers
-        networkLayers.layers = networkLayers.layers.filter(layer => layer.id !== nodeId);
-        networkLayers.connections = networkLayers.connections.filter(conn => 
-            conn.source !== nodeId && conn.target !== nodeId
-        );
-        
-        // Remove the node
-        node.parentNode.removeChild(node);
-        
-        // Update layer connectivity
-        updateLayerConnectivity();
-    }
-    
-    // Open layer editor modal
-    function openLayerEditor(node) {
-        if (!node) return;
-        
-        const nodeId = node.getAttribute('data-id');
-        const nodeType = node.getAttribute('data-type');
-        const nodeName = node.getAttribute('data-name');
-        const dimensions = node.getAttribute('data-dimensions');
-        
-        // Trigger custom event
-        const event = new CustomEvent('openLayerEditor', { 
-            detail: { id: nodeId, type: nodeType, name: nodeName, dimensions: dimensions }
-        });
-        document.dispatchEvent(event);
-    }
-    
-    // Update connections when nodes are moved
-    function updateConnections() {
-        const connections = document.querySelectorAll('.connection');
-        connections.forEach(connection => {
-            const sourceId = connection.getAttribute('data-source');
-            const targetId = connection.getAttribute('data-target');
-            
-            const sourceNode = document.querySelector(`.canvas-node[data-id="${sourceId}"]`);
-            const targetNode = document.querySelector(`.canvas-node[data-id="${targetId}"]`);
-            
-            if (sourceNode && targetNode) {
-                const sourcePort = sourceNode.querySelector('.port-out');
-                const targetPort = targetNode.querySelector('.port-in');
-                
-                if (sourcePort && targetPort) {
-                    const sourceRect = sourcePort.getBoundingClientRect();
-                    const targetRect = targetPort.getBoundingClientRect();
-                    const canvasRect = canvas.getBoundingClientRect();
-                    
-                    const startX = sourceRect.left + sourceRect.width / 2 - canvasRect.left;
-                    const startY = sourceRect.top + sourceRect.height / 2 - canvasRect.top;
-                    const endX = targetRect.left + targetRect.width / 2 - canvasRect.left;
-                    const endY = targetRect.top + targetRect.height / 2 - canvasRect.top;
-                    
-                    const length = Math.sqrt(Math.pow(endX - startX, 2) + Math.pow(endY - startY, 2));
-                    const angle = Math.atan2(endY - startY, endX - startX) * 180 / Math.PI;
-                    
-                    connection.style.left = `${startX}px`;
-                    connection.style.top = `${startY}px`;
-                    connection.style.width = `${length}px`;
-                    connection.style.transform = `rotate(${angle}deg)`;
-                }
-            } else {
-                // If either node is missing, remove the connection
-                if (connection.parentNode) {
-                    connection.parentNode.removeChild(connection);
-                    
-                    // Remove from the connections array
-                    const connIndex = networkLayers.connections.findIndex(conn => 
-                        conn.source === sourceId && conn.target === targetId
-                    );
-                    if (connIndex !== -1) {
-                        networkLayers.connections.splice(connIndex, 1);
-                    }
-                }
-            }
-        });
-    }
-    
-    // Get the current network architecture
-    function getNetworkArchitecture() {
-        return networkLayers;
-    }
-    
-    // Clear all nodes from the canvas
-    function clearAllNodes() {
-        // Clear all nodes and connections
-        document.querySelectorAll('.canvas-node, .connection').forEach(el => {
-            el.parentNode.removeChild(el);
-        });
-        
-        // Reset network layers
-        networkLayers = {
-            layers: [],
-            connections: []
-        };
-        
-        // Reset layer counter
-        window.neuralNetwork.resetLayerCounter();
-        
-        // Show the canvas hint
-        const canvasHint = document.querySelector('.canvas-hint');
-        if (canvasHint) {
-            canvasHint.style.display = 'block';
-        }
-        
-        // Trigger network updated event
-        const event = new CustomEvent('networkUpdated', { detail: networkLayers });
-        document.dispatchEvent(event);
-    }
-    
-    // Export functions
-    window.dragDrop = {
-        getNetworkArchitecture,
-        clearAllNodes,
-        updateConnections
-    };
-} 
+            const pathId = `connection-${sourceId}-${targetId}`

js/main.js

@@ -405,7 +405,20 @@ document.addEventListener('DOMContentLoaded', () => {
                 
             case 'conv':
                 // Convolutional layer parameters
+                // Get input and output shapes - may be calculated or null at first
+                const inputShape = layerConfig.inputShape || ['?', '?', '?'];
+                const outputShape = layerConfig.outputShape || ['?', '?', layerConfig.filters];
+                
                 layerForm.innerHTML += `
+                    <div class="form-group">
+                        <label>Input Shape:</label>
+                        <div class="form-row">
+                            <input type="number" id="conv-input-h" min="1" value="${inputShape[0] === '?' ? 28 : inputShape[0]}" placeholder="Height">
+                            <input type="number" id="conv-input-w" min="1" value="${inputShape[1] === '?' ? 28 : inputShape[1]}" placeholder="Width">
+                            <input type="number" id="conv-input-c" min="1" value="${inputShape[2] === '?' ? 1 : inputShape[2]}" placeholder="Channels">
+                        </div>
+                        <small>Input dimensions: H × W × C</small>
+                    </div>
                     <div class="form-group">
                         <label>Filters:</label>
                         <input type="number" id="conv-filters" min="1" value="${layerConfig.filters}" placeholder="Number of filters">
@@ -442,12 +455,92 @@ document.addEventListener('DOMContentLoaded', () => {
                             <option value="leaky_relu" ${layerConfig.activation === 'leaky_relu' ? 'selected' : ''}>Leaky ReLU</option>
                         </select>
                     </div>
+                    <div class="form-group">
+                        <label>Output Shape (calculated):</label>
+                        <div class="output-shape-display" id="conv-output-shape">
+                            [${outputShape.join(' × ')}]
+                        </div>
+                        <small>Output dimensions: H × W × Filters</small>
+                    </div>
+                    <div class="form-group">
+                        <label>Parameters (calculated):</label>
+                        <div class="parameters-display" id="conv-parameters">
+                            Calculating...
+                        </div>
+                    </div>
                 `;
+                
+                // Add event listeners to calculate output shape and parameters in real-time
+                setTimeout(() => {
+                    const inputH = document.getElementById('conv-input-h');
+                    const inputW = document.getElementById('conv-input-w');
+                    const inputC = document.getElementById('conv-input-c');
+                    const filters = document.getElementById('conv-filters');
+                    const kernelH = document.getElementById('kernel-size-h');
+                    const kernelW = document.getElementById('kernel-size-w');
+                    const strideH = document.getElementById('stride-h');
+                    const strideW = document.getElementById('stride-w');
+                    const paddingType = document.getElementById('padding-type');
+                    const outputShapeDisplay = document.getElementById('conv-output-shape');
+                    const parametersDisplay = document.getElementById('conv-parameters');
+                    
+                    const updateOutputShape = () => {
+                        const h = parseInt(inputH.value);
+                        const w = parseInt(inputW.value);
+                        const c = parseInt(inputC.value);
+                        const f = parseInt(filters.value);
+                        const kh = parseInt(kernelH.value);
+                        const kw = parseInt(kernelW.value);
+                        const sh = parseInt(strideH.value);
+                        const sw = parseInt(strideW.value);
+                        const padding = paddingType.value;
+                        
+                        // Calculate output dimensions
+                        const pH = padding === 'same' ? Math.floor(kh / 2) : 0;
+                        const pW = padding === 'same' ? Math.floor(kw / 2) : 0;
+                        
+                        const outH = Math.floor((h - kh + 2 * pH) / sh) + 1;
+                        const outW = Math.floor((w - kw + 2 * pW) / sw) + 1;
+                        
+                        // Update output shape display
+                        outputShapeDisplay.textContent = `[${outH} × ${outW} × ${f}]`;
+                        
+                        // Calculate parameters
+                        const params = kh * kw * c * f + f; // weights + bias
+                        parametersDisplay.textContent = formatNumber(params);
+                        
+                        // Store for saving
+                        layerConfig.inputShape = [h, w, c];
+                        layerConfig.outputShape = [outH, outW, f];
+                        layerConfig.parameters = params;
+                    };
+                    
+                    // Attach event listeners to all inputs
+                    [inputH, inputW, inputC, filters, kernelH, kernelW, strideH, strideW, paddingType].forEach(
+                        input => input.addEventListener('input', updateOutputShape)
+                    );
+                    
+                    // Initialize values
+                    updateOutputShape();
+                }, 100);
                 break;
                 
             case 'pool':
                 // Pooling layer parameters
+                // Get input and output shapes
+                const poolInputShape = layerConfig.inputShape || ['?', '?', '?'];
+                const poolOutputShape = layerConfig.outputShape || ['?', '?', '?'];
+                
                 layerForm.innerHTML += `
+                    <div class="form-group">
+                        <label>Input Shape:</label>
+                        <div class="form-row">
+                            <input type="number" id="pool-input-h" min="1" value="${poolInputShape[0] === '?' ? 28 : poolInputShape[0]}" placeholder="Height">
+                            <input type="number" id="pool-input-w" min="1" value="${poolInputShape[1] === '?' ? 28 : poolInputShape[1]}" placeholder="Width">
+                            <input type="number" id="pool-input-c" min="1" value="${poolInputShape[2] === '?' ? 1 : poolInputShape[2]}" placeholder="Channels">
+                        </div>
+                        <small>Input dimensions: H × W × C</small>
+                    </div>
                     <div class="form-group">
                         <label>Pool Size:</label>
                         <div class="form-row">
@@ -476,7 +569,61 @@ document.addEventListener('DOMContentLoaded', () => {
                             <option value="avg">Average Pooling</option>
                         </select>
                     </div>
+                    <div class="form-group">
+                        <label>Output Shape (calculated):</label>
+                        <div class="output-shape-display" id="pool-output-shape">
+                            [${poolOutputShape.join(' × ')}]
+                        </div>
+                        <small>Output dimensions: H × W × C</small>
+                    </div>
                 `;
+                
+                // Add event listeners to calculate output shape in real-time
+                setTimeout(() => {
+                    const inputH = document.getElementById('pool-input-h');
+                    const inputW = document.getElementById('pool-input-w');
+                    const inputC = document.getElementById('pool-input-c');
+                    const poolH = document.getElementById('pool-size-h');
+                    const poolW = document.getElementById('pool-size-w');
+                    const strideH = document.getElementById('pool-stride-h');
+                    const strideW = document.getElementById('pool-stride-w');
+                    const paddingType = document.getElementById('pool-padding');
+                    const outputShapeDisplay = document.getElementById('pool-output-shape');
+                    
+                    const updateOutputShape = () => {
+                        const h = parseInt(inputH.value);
+                        const w = parseInt(inputW.value);
+                        const c = parseInt(inputC.value);
+                        const ph = parseInt(poolH.value);
+                        const pw = parseInt(poolW.value);
+                        const sh = parseInt(strideH.value);
+                        const sw = parseInt(strideW.value);
+                        const padding = paddingType.value;
+                        
+                        // Calculate output dimensions
+                        const padH = padding === 'same' ? Math.floor(ph / 2) : 0;
+                        const padW = padding === 'same' ? Math.floor(pw / 2) : 0;
+                        
+                        const outH = Math.floor((h - ph + 2 * padH) / sh) + 1;
+                        const outW = Math.floor((w - pw + 2 * padW) / sw) + 1;
+                        
+                        // Update output shape display
+                        outputShapeDisplay.textContent = `[${outH} × ${outW} × ${c}]`;
+                        
+                        // Store for saving
+                        layerConfig.inputShape = [h, w, c];
+                        layerConfig.outputShape = [outH, outW, c];
+                        layerConfig.parameters = 0; // Pooling has no parameters
+                    };
+                    
+                    // Attach event listeners to all inputs
+                    [inputH, inputW, inputC, poolH, poolW, strideH, strideW, paddingType].forEach(
+                        input => input.addEventListener('input', updateOutputShape)
+                    );
+                    
+                    // Initialize values
+                    updateOutputShape();
+                }, 100);
                 break;
                 
             case 'linear':
@@ -591,28 +738,169 @@ document.addEventListener('DOMContentLoaded', () => {
         const values = {};
         const inputs = form.querySelectorAll('input, select');
         inputs.forEach(input => {
-            values[input.id] = input.value;
+            if (input.type === 'checkbox') {
+                values[input.id] = input.checked;
+            } else {
+                values[input.id] = input.value;
+            }
         });
         
         // Update node configuration
-        node.layerConfig = {
-            type: nodeType,
-            shape: [
-                parseInt(values['input-height']),
-                parseInt(values['input-width']),
-                parseInt(values['input-channels'])
-            ],
-            batchSize: parseInt(values['batch-size']),
-            units: parseInt(values['hidden-units']),
-            activation: values['hidden-activation'],
-            dropoutRate: parseFloat(values['dropout-rate']),
-            useBias: values['use-bias'] === 'true',
-            learningRate: parseFloat(values['learning-rate-slider']),
-            lossFunction: values['loss-function'],
-            optimizer: values['optimizer'],
-            inputFeatures: parseInt(values['input-features']),
-            outputFeatures: parseInt(values['output-features'])
-        };
+        node.layerConfig = node.layerConfig || {};
+        const layerConfig = node.layerConfig;
+        
+        switch (nodeType) {
+            case 'input':
+                layerConfig.shape = [
+                    parseInt(values['input-height']) || 28,
+                    parseInt(values['input-width']) || 28,
+                    parseInt(values['input-channels']) || 1
+                ];
+                layerConfig.batchSize = parseInt(values['batch-size']) || 32;
+                layerConfig.outputShape = layerConfig.shape;
+                layerConfig.parameters = 0;
+                break;
+                
+            case 'hidden':
+                layerConfig.units = parseInt(values['hidden-units']) || 128;
+                layerConfig.activation = values['hidden-activation'] || 'relu';
+                layerConfig.dropoutRate = parseFloat(values['dropout-rate']) || 0.2;
+                layerConfig.useBias = values['use-bias'] === true;
+                layerConfig.outputShape = [layerConfig.units];
+                
+                // Calculate parameters if input shape is available
+                if (layerConfig.inputShape) {
+                    const inputUnits = Array.isArray(layerConfig.inputShape) ? 
+                        layerConfig.inputShape.reduce((a, b) => a * b, 1) : layerConfig.inputShape;
+                    layerConfig.parameters = (inputUnits * layerConfig.units) + (layerConfig.useBias ? layerConfig.units : 0);
+                }
+                break;
+                
+            case 'output':
+                layerConfig.units = parseInt(values['output-units']) || 10;
+                layerConfig.activation = values['output-activation'] || 'softmax';
+                layerConfig.useBias = values['output-use-bias'] === true;
+                layerConfig.outputShape = [layerConfig.units];
+                
+                // Calculate parameters if input shape is available
+                if (layerConfig.inputShape) {
+                    const inputUnits = Array.isArray(layerConfig.inputShape) ? 
+                        layerConfig.inputShape.reduce((a, b) => a * b, 1) : layerConfig.inputShape;
+                    layerConfig.parameters = (inputUnits * layerConfig.units) + (layerConfig.useBias ? layerConfig.units : 0);
+                }
+                break;
+                
+            case 'conv':
+                // Process input shape if available in form
+                if (values['conv-input-h'] && values['conv-input-w'] && values['conv-input-c']) {
+                    layerConfig.inputShape = [
+                        parseInt(values['conv-input-h']) || 28,
+                        parseInt(values['conv-input-w']) || 28,
+                        parseInt(values['conv-input-c']) || 1
+                    ];
+                }
+                
+                // Process configuration
+                layerConfig.filters = parseInt(values['conv-filters']) || 32;
+                layerConfig.kernelSize = [
+                    parseInt(values['kernel-size-h']) || 3,
+                    parseInt(values['kernel-size-w']) || 3
+                ];
+                layerConfig.strides = [
+                    parseInt(values['stride-h']) || 1,
+                    parseInt(values['stride-w']) || 1
+                ];
+                layerConfig.padding = values['padding-type'] || 'valid';
+                layerConfig.activation = values['conv-activation'] || 'relu';
+                layerConfig.useBias = true; // Default to true for CNN
+                
+                // Calculate output shape if input shape is available
+                if (layerConfig.inputShape) {
+                    const padding = layerConfig.padding === 'same' ? 
+                        Math.floor(layerConfig.kernelSize[0] / 2) : 0;
+                        
+                    const outH = Math.floor(
+                        (layerConfig.inputShape[0] - layerConfig.kernelSize[0] + 2 * padding) / 
+                        layerConfig.strides[0]
+                    ) + 1;
+                    
+                    const outW = Math.floor(
+                        (layerConfig.inputShape[1] - layerConfig.kernelSize[1] + 2 * padding) / 
+                        layerConfig.strides[1]
+                    ) + 1;
+                    
+                    layerConfig.outputShape = [outH, outW, layerConfig.filters];
+                    
+                    // Calculate parameters
+                    const kernelParams = layerConfig.kernelSize[0] * layerConfig.kernelSize[1] * 
+                                        layerConfig.inputShape[2] * layerConfig.filters;
+                    const biasParams = layerConfig.filters;
+                    layerConfig.parameters = kernelParams + biasParams;
+                }
+                break;
+                
+            case 'pool':
+                // Process input shape if available in form
+                if (values['pool-input-h'] && values['pool-input-w'] && values['pool-input-c']) {
+                    layerConfig.inputShape = [
+                        parseInt(values['pool-input-h']) || 28,
+                        parseInt(values['pool-input-w']) || 28,
+                        parseInt(values['pool-input-c']) || 1
+                    ];
+                }
+                
+                // Process configuration
+                layerConfig.poolSize = [
+                    parseInt(values['pool-size-h']) || 2,
+                    parseInt(values['pool-size-w']) || 2
+                ];
+                layerConfig.strides = [
+                    parseInt(values['pool-stride-h']) || 2,
+                    parseInt(values['pool-stride-w']) || 2
+                ];
+                layerConfig.padding = values['pool-padding'] || 'valid';
+                layerConfig.poolType = values['pool-type'] || 'max';
+                
+                // Calculate output shape if input shape is available
+                if (layerConfig.inputShape) {
+                    const poolPadding = layerConfig.padding === 'same' ? 
+                        Math.floor(layerConfig.poolSize[0] / 2) : 0;
+                        
+                    const poolOutH = Math.floor(
+                        (layerConfig.inputShape[0] - layerConfig.poolSize[0] + 2 * poolPadding) / 
+                        layerConfig.strides[0]
+                    ) + 1;
+                    
+                    const poolOutW = Math.floor(
+                        (layerConfig.inputShape[1] - layerConfig.poolSize[1] + 2 * poolPadding) / 
+                        layerConfig.strides[1]
+                    ) + 1;
+                    
+                    layerConfig.outputShape = [poolOutH, poolOutW, layerConfig.inputShape[2]];
+                }
+                
+                // Pooling has no parameters
+                layerConfig.parameters = 0;
+                break;
+                
+            case 'linear':
+                layerConfig.inputFeatures = parseInt(values['input-features']) || 1;
+                layerConfig.outputFeatures = parseInt(values['output-features']) || 1;
+                layerConfig.useBias = values['linear-use-bias'] === true;
+                layerConfig.learningRate = parseFloat(values['learning-rate-slider']) || 0.01;
+                layerConfig.activation = values['linear-activation'] || 'linear';
+                layerConfig.optimizer = values['optimizer'] || 'sgd';
+                layerConfig.lossFunction = values['loss-function'] || 'mse';
+                layerConfig.inputShape = [layerConfig.inputFeatures];
+                layerConfig.outputShape = [layerConfig.outputFeatures];
+                
+                // Calculate parameters
+                layerConfig.parameters = layerConfig.inputFeatures * layerConfig.outputFeatures;
+                if (layerConfig.useBias) {
+                    layerConfig.parameters += layerConfig.outputFeatures;
+                }
+                break;
+        }
         
         // Update node title
         const nodeTitle = node.querySelector('.node-title');
@@ -623,37 +911,55 @@ document.addEventListener('DOMContentLoaded', () => {
         // Update node data attribute
         node.setAttribute('data-name', nodeType.charAt(0).toUpperCase() + nodeType.slice(1));
         
-        // Update dimensions based on layer type
+        // Update dimensions and parameter display based on layer type
         let dimensions = '';
         switch (nodeType) {
             case 'input':
-                dimensions = values['input-height'] + ' × ' + values['input-width'] + ' × ' + values['input-channels'];
+                dimensions = layerConfig.shape.join(' × ');
                 break;
                 
             case 'hidden':
             case 'output':
-                dimensions = values['hidden-units'];
+                dimensions = layerConfig.units.toString();
                 break;
                 
             case 'conv':
-                dimensions = values['conv-filters'] + ' × ' + values['kernel-size-h'] + ' × ' + values['kernel-size-w'];
+                if (layerConfig.inputShape && layerConfig.outputShape) {
+                    // Show input -> output shape transformation
+                    dimensions = `${layerConfig.inputShape[0]}×${layerConfig.inputShape[1]}×${layerConfig.inputShape[2]} → ${layerConfig.outputShape[0]}×${layerConfig.outputShape[1]}×${layerConfig.outputShape[2]}`;
+                } else {
+                    dimensions = `? → ${layerConfig.filters} filters`;
+                }
                 break;
                 
             case 'pool':
-                dimensions = values['pool-size-h'] + ' × ' + values['pool-size-w'];
+                if (layerConfig.inputShape && layerConfig.outputShape) {
+                    // Show input -> output shape transformation
+                    dimensions = `${layerConfig.inputShape[0]}×${layerConfig.inputShape[1]}×${layerConfig.inputShape[2]} → ${layerConfig.outputShape[0]}×${layerConfig.outputShape[1]}×${layerConfig.outputShape[2]}`;
+                } else {
+                    dimensions = `? → ?`;
+                }
                 break;
                 
             case 'linear':
-                dimensions = values['input-features'] + ' → ' + values['output-features'];
+                dimensions = `${layerConfig.inputFeatures} → ${layerConfig.outputFeatures}`;
                 break;
         }
         
-        // Update node dimensions
+        // Update node dimensions display
         const nodeDimensions = node.querySelector('.node-dimensions');
         if (nodeDimensions) {
             nodeDimensions.textContent = dimensions;
         }
         
+        // Update parameters display if available
+        const nodeParameters = node.querySelector('.node-parameters');
+        if (nodeParameters && layerConfig.parameters !== undefined) {
+            nodeParameters.textContent = `Params: ${formatNumber(layerConfig.parameters)}`;
+        } else if (nodeParameters) {
+            nodeParameters.textContent = 'Params: ?';
+        }
+        
         // Update node data attribute
         node.setAttribute('data-dimensions', dimensions);
         
@@ -664,11 +970,96 @@ document.addEventListener('DOMContentLoaded', () => {
         if (layerIndex !== -1) {
             networkLayers.layers[layerIndex].name = nodeType.charAt(0).toUpperCase() + nodeType.slice(1);
             networkLayers.layers[layerIndex].dimensions = dimensions;
+            networkLayers.layers[layerIndex].config = layerConfig;
+            
+            // Add parameter count to the layer
+            networkLayers.layers[layerIndex].parameters = layerConfig.parameters;
         }
         
         // Trigger network updated event
         const event = new CustomEvent('networkUpdated', { detail: networkLayers });
         document.dispatchEvent(event);
+        
+        // Update connected nodes to propagate shape changes
+        updateNodeConnections(node, layerId);
+    }
+    
+    // Helper function to update connections between nodes when shapes change
+    function updateNodeConnections(sourceNode, sourceId) {
+        // Find all connections from this source node
+        const connections = document.querySelectorAll(`.connection[data-source="${sourceId}"]`);
+        
+        connections.forEach(connection => {
+            const targetId = connection.getAttribute('data-target');
+            const targetNode = document.querySelector(`.canvas-node[data-id="${targetId}"]`);
+            
+            if (targetNode && sourceNode.layerConfig && sourceNode.layerConfig.outputShape) {
+                // Update target node with source node's output shape as its input shape
+                if (!targetNode.layerConfig) {
+                    targetNode.layerConfig = {};
+                }
+                
+                targetNode.layerConfig.inputShape = sourceNode.layerConfig.outputShape;
+                
+                // Update parameter calculation
+                window.neuralNetwork.calculateParameters(
+                    targetNode.getAttribute('data-type'),
+                    targetNode.layerConfig,
+                    sourceNode.layerConfig
+                );
+                
+                // Update display
+                updateNodeDisplay(targetNode);
+                
+                // Recursively update downstream nodes
+                updateNodeConnections(targetNode, targetId);
+            }
+        });
+    }
+    
+    // Helper function to update a node's display
+    function updateNodeDisplay(node) {
+        if (!node || !node.layerConfig) return;
+        
+        const nodeType = node.getAttribute('data-type');
+        const layerConfig = node.layerConfig;
+        
+        // Create dimensions string
+        let dimensions = '';
+        switch (nodeType) {
+            case 'conv':
+            case 'pool':
+                if (layerConfig.inputShape && layerConfig.outputShape) {
+                    dimensions = `${layerConfig.inputShape[0]}×${layerConfig.inputShape[1]}×${layerConfig.inputShape[2]} → ${layerConfig.outputShape[0]}×${layerConfig.outputShape[1]}×${layerConfig.outputShape[2]}`;
+                }
+                break;
+                
+            case 'hidden':
+            case 'output':
+                dimensions = layerConfig.units.toString();
+                break;
+                
+            case 'linear':
+                dimensions = `${layerConfig.inputFeatures} → ${layerConfig.outputFeatures}`;
+                break;
+        }
+        
+        // Update dimensions display
+        if (dimensions) {
+            const nodeDimensions = node.querySelector('.node-dimensions');
+            if (nodeDimensions) {
+                nodeDimensions.textContent = dimensions;
+                node.setAttribute('data-dimensions', dimensions);
+            }
+        }
+        
+        // Update parameters display
+        if (layerConfig.parameters !== undefined) {
+            const nodeParameters = node.querySelector('.node-parameters');
+            if (nodeParameters) {
+                nodeParameters.textContent = `Params: ${formatNumber(layerConfig.parameters)}`;
+            }
+        }
     }
     
     // Handle sample selection

js/neural-network.js

@@ -153,10 +153,10 @@
     }
     
     /**
-     * Calculate the number of parameters for a layer
+     * Calculate parameters for a layer
      * @param {string} layerType - The type of the layer
      * @param {Object} config - Layer configuration
-     * @param {Object} prevLayerConfig - Previous layer configuration (for connections)
+     * @param {Object} prevLayerConfig - Configuration of the previous connected layer
      * @returns {number} - Number of trainable parameters
      */
     function calculateParameters(layerType, config, prevLayerConfig = null) {
@@ -169,10 +169,17 @@
                 
             case 'hidden':
                 if (prevLayerConfig) {
-                    const inputUnits = prevLayerConfig.units || 
-                                      (prevLayerConfig.shape ? 
-                                       prevLayerConfig.shape.reduce((a, b) => a * b, 1) : 
-                                       784);
+                    // Calculate input units from previous layer shape or units
+                    let inputUnits;
+                    if (prevLayerConfig.outputShape && Array.isArray(prevLayerConfig.outputShape)) {
+                        inputUnits = prevLayerConfig.outputShape.reduce((a, b) => a * b, 1);
+                    } else if (prevLayerConfig.units) {
+                        inputUnits = prevLayerConfig.units;
+                    } else if (prevLayerConfig.shape) {
+                        inputUnits = prevLayerConfig.shape.reduce((a, b) => a * b, 1);
+                    } else {
+                        inputUnits = 784; // Default fallback
+                    }
                     
                     // Weight parameters: input_units * output_units
                     parameters = inputUnits * config.units;
@@ -186,7 +193,15 @@
                 
             case 'output':
                 if (prevLayerConfig) {
-                    const inputUnits = prevLayerConfig.units || 128;
+                    // Calculate input units from previous layer
+                    let inputUnits;
+                    if (prevLayerConfig.outputShape && Array.isArray(prevLayerConfig.outputShape)) {
+                        inputUnits = prevLayerConfig.outputShape.reduce((a, b) => a * b, 1);
+                    } else if (prevLayerConfig.units) {
+                        inputUnits = prevLayerConfig.units;
+                    } else {
+                        inputUnits = 128; // Default fallback
+                    }
                     
                     // Weight parameters: input_units * output_units
                     parameters = inputUnits * config.units;
@@ -200,9 +215,17 @@
                 
             case 'conv':
                 if (prevLayerConfig) {
-                    const inputChannels = prevLayerConfig.shape ? 
-                                         prevLayerConfig.shape[2] || 1 : 
-                                         (prevLayerConfig.filters || 1);
+                    // Get input channels from previous layer
+                    let inputChannels;
+                    if (prevLayerConfig.outputShape && prevLayerConfig.outputShape.length > 2) {
+                        inputChannels = prevLayerConfig.outputShape[2];
+                    } else if (prevLayerConfig.shape && prevLayerConfig.shape.length > 2) {
+                        inputChannels = prevLayerConfig.shape[2];
+                    } else if (prevLayerConfig.filters) {
+                        inputChannels = prevLayerConfig.filters;
+                    } else {
+                        inputChannels = 1; // Default fallback
+                    }
                     
                     // Weight parameters: kernel_height * kernel_width * input_channels * filters
                     const kernelSize = Array.isArray(config.kernelSize) ? 
@@ -215,11 +238,31 @@
                     if (config.useBias) {
                         parameters += config.filters;
                     }
+                    
+                    // Calculate and store output shape
+                    if (prevLayerConfig.shape || prevLayerConfig.outputShape) {
+                        const inputShape = prevLayerConfig.outputShape || prevLayerConfig.shape;
+                        const padding = config.padding === 'same' ? Math.floor(config.kernelSize[0] / 2) : 0;
+                        const outputHeight = Math.floor((inputShape[0] - config.kernelSize[0] + 2 * padding) / config.strides[0]) + 1;
+                        const outputWidth = Math.floor((inputShape[1] - config.kernelSize[1] + 2 * padding) / config.strides[1]) + 1;
+                        
+                        config.outputShape = [outputHeight, outputWidth, config.filters];
+                    }
                 }
                 break;
                 
             case 'pool':
                 parameters = 0; // Pooling layers have no trainable parameters
+                
+                // Calculate and store output shape
+                if (prevLayerConfig && (prevLayerConfig.shape || prevLayerConfig.outputShape)) {
+                    const inputShape = prevLayerConfig.outputShape || prevLayerConfig.shape;
+                    const padding = config.padding === 'same' ? Math.floor(config.poolSize[0] / 2) : 0;
+                    const outputHeight = Math.floor((inputShape[0] - config.poolSize[0] + 2 * padding) / config.strides[0]) + 1;
+                    const outputWidth = Math.floor((inputShape[1] - config.poolSize[1] + 2 * padding) / config.strides[1]) + 1;
+                    
+                    config.outputShape = [outputHeight, outputWidth, inputShape[2]];
+                }
                 break;
                 
             default: