shape_inference.py

"""
Automatic Shape Inference and Resolution Pass

This pass automatically infers tensor shapes throughout the MIR graph and
inserts necessary reshape/flatten operations to ensure compatibility between
layers (e.g., automatic flatten before Linear layers).

Key Features:
- Forward shape propagation through entire graph
- Automatic flatten insertion before fully connected layers
- Dynamic shape computation for pooling and convolution
- Shape validation and error detection

This eliminates manual shape handling and prevents common shape mismatch errors.
"""

import numpy as np
from typing import Dict, List, Tuple, Any, Optional


class ShapeInferencePass:
    """
    Automatic shape inference and resolution for MIR graphs.
    
    This pass:
    1. Propagates shapes forward through the graph
    2. Inserts flatten operations where needed
    3. Validates shape compatibility
    4. Annotates each node with output shape metadata
    """
    
    def __init__(self):
        self.shape_map = {}  # node_id -> output_shape
        self.inserted_nodes = []  # Track inserted flatten nodes
    
    def run(self, mir: Dict[str, Any]) -> Dict[str, Any]:
        """
        Run shape inference pass on MIR graph.
        
        Args:
            mir: MIR graph dictionary
        
        Returns:
            mir: Updated MIR with shape annotations and inserted flatten ops
        """
        print("\n🔍 Running Shape Inference Pass...")
        
        # Get input shape
        if 'input_shape' in mir:
            input_shape = mir['input_shape']
        elif 'graph' in mir and 'input_shape' in mir['graph']:
            input_shape = mir['graph']['input_shape']
        else:
            raise ValueError("Cannot find input_shape in MIR")
        
        # Initialize shape map with input
        self.shape_map['input'] = input_shape
        
        # Get nodes
        if 'nodes' in mir:
            nodes = mir['nodes']
        elif 'graph' in mir and 'nodes' in mir['graph']:
            nodes = mir['graph']['nodes']
        else:
            raise ValueError("Cannot find nodes in MIR")
        
        # Propagate shapes through graph
        modified_nodes = []
        for i, node in enumerate(nodes):
            # Infer input shapes
            input_shapes = self._get_input_shapes(node)
            
            # Check if reshape/flatten needed before this node
            if node['op_type'] == 'linear':
                # Linear expects 2D input [batch, features]
                if input_shapes and len(input_shapes[0]) > 2:
                    # Need to insert flatten
                    flatten_node = self._create_flatten_node(node, i)
                    modified_nodes.append(flatten_node)
                    
                    # Update input shape for linear layer
                    input_shapes = [self._infer_flatten_shape(input_shapes[0])]
                    
                    # Update node to use flatten output as input
                    original_input = node['inputs'][0]
                    node['inputs'][0] = flatten_node['id']
                    
                    print(f"   ✓ Inserted flatten before '{node['id']}' (Linear layer)")
            
            # Add node
            modified_nodes.append(node)
            
            # Infer output shape for this node
            output_shape = self._infer_node_output_shape(
                node['op_type'],
                input_shapes,
                node
            )
            
            # Store in shape map
            self.shape_map[node['id']] = output_shape
            
            # Annotate node with shape metadata
            node['shape_info'] = {
                'input_shapes': input_shapes,
                'output_shape': output_shape
            }
        
        # Update MIR with modified nodes
        if 'nodes' in mir:
            mir['nodes'] = modified_nodes
        elif 'graph' in mir:
            mir['graph']['nodes'] = modified_nodes
        
        # Add shape map to metadata
        if 'metadata' not in mir:
            mir['metadata'] = {}
        mir['metadata']['shape_map'] = self.shape_map
        mir['metadata']['shape_inference_applied'] = True
        
        print(f"✅ Shape inference complete")
        print(f"   - Total nodes: {len(modified_nodes)}")
        print(f"   - Flatten ops inserted: {len(self.inserted_nodes)}")
        print(f"   - Shape map entries: {len(self.shape_map)}")
        
        return mir
    
    def _get_input_shapes(self, node: Dict[str, Any]) -> List[List[int]]:
        """Get input shapes for a node from the shape map."""
        input_shapes = []
        
        for input_id in node.get('inputs', []):
            if input_id in self.shape_map:
                input_shapes.append(self.shape_map[input_id])
            else:
                # Unknown shape - will be inferred later
                input_shapes.append(None)
        
        return input_shapes
    
    def _create_flatten_node(self, target_node: Dict[str, Any], insert_index: int) -> Dict[str, Any]:
        """Create a flatten node to insert before a Linear layer."""
        original_input = target_node['inputs'][0]
        flatten_id = f"flatten_auto_{insert_index}"
        
        flatten_node = {
            'id': flatten_id,
            'op_type': 'flatten',
            'inputs': [original_input],
            'outputs': [flatten_id],
            'params': {
                'start_dim': 1  # Flatten all dims except batch
            },
            'metadata': {
                'auto_inserted': True,
                'inserted_before': target_node['id']
            }
        }
        
        self.inserted_nodes.append(flatten_id)
        return flatten_node
    
    def _infer_flatten_shape(self, input_shape: List[int]) -> List[int]:
        """Infer output shape of flatten operation."""
        if len(input_shape) <= 2:
            return input_shape
        
        # Flatten all dims except batch: [batch, ...] -> [batch, features]
        batch_size = input_shape[0]
        features = int(np.prod(input_shape[1:]))
        
        return [batch_size, features]
    
    def _infer_node_output_shape(
        self,
        op_type: str,
        input_shapes: List[List[int]],
        node: Dict[str, Any]
    ) -> List[int]:
        """
        Infer output shape for a node based on operation type.
        
        Args:
            op_type: Type of operation (conv2d, linear, etc.)
            input_shapes: List of input shapes
            node: Node dictionary with parameters
        
        Returns:
            output_shape: Inferred output shape
        """
        if not input_shapes or input_shapes[0] is None:
            return None
        
        input_shape = input_shapes[0]
        
        if op_type == 'conv2d':
            return self._infer_conv2d_shape(input_shape, node)
        
        elif op_type == 'linear':
            return self._infer_linear_shape(input_shape, node)
        
        elif op_type == 'relu':
            # ReLU doesn't change shape
            return input_shape
        
        elif op_type == 'batchnorm2d':
            # BatchNorm doesn't change shape
            return input_shape
        
        elif op_type in ['maxpool2d', 'max_pool2d']:
            return self._infer_maxpool2d_shape(input_shape, node)
        
        elif op_type in ['adaptiveavgpool2d', 'adaptive_avg_pool2d']:
            return self._infer_adaptive_avgpool2d_shape(input_shape, node)
        
        elif op_type == 'flatten':
            return self._infer_flatten_shape(input_shape)
        
        elif op_type == 'reshape':
            return self._infer_reshape_shape(input_shape, node)
        
        elif op_type in ['add', 'concat', 'concatenate']:
            # Add/concat keep first input shape (simplified)
            return input_shape
        
        else:
            # Unknown op - return input shape as fallback
            print(f"   ⚠️  Unknown op_type '{op_type}' - assuming shape unchanged")
            return input_shape
    
    def _infer_conv2d_shape(self, input_shape: List[int], node: Dict[str, Any]) -> List[int]:
        """Infer Conv2D output shape."""
        params = node.get('params', node.get('parameters', {}))
        
        batch = input_shape[0]
        in_channels = input_shape[1]
        in_h = input_shape[2]
        in_w = input_shape[3]
        
        # Get conv parameters
        out_channels = params.get('out_channels', in_channels)
        kernel_size = params.get('kernel_size', [3, 3])
        stride = params.get('stride', [1, 1])
        padding = params.get('padding', [0, 0])
        
        # Handle scalar kernel_size/stride/padding
        if isinstance(kernel_size, int):
            kernel_size = [kernel_size, kernel_size]
        if isinstance(stride, int):
            stride = [stride, stride]
        if isinstance(padding, int):
            padding = [padding, padding]
        
        # Calculate output dimensions
        out_h = (in_h + 2 * padding[0] - kernel_size[0]) // stride[0] + 1
        out_w = (in_w + 2 * padding[1] - kernel_size[1]) // stride[1] + 1
        
        return [batch, out_channels, out_h, out_w]
    
    def _infer_linear_shape(self, input_shape: List[int], node: Dict[str, Any]) -> List[int]:
        """Infer Linear/FC output shape."""
        params = node.get('params', node.get('parameters', {}))
        
        batch = input_shape[0]
        out_features = params.get('out', params.get('out_features', 10))
        
        return [batch, out_features]
    
    def _infer_maxpool2d_shape(self, input_shape: List[int], node: Dict[str, Any]) -> List[int]:
        """Infer MaxPool2D output shape."""
        params = node.get('params', node.get('parameters', {}))
        
        batch = input_shape[0]
        channels = input_shape[1]
        in_h = input_shape[2]
        in_w = input_shape[3]
        
        # Get pooling parameters
        kernel_size = params.get('kernel_size', [2, 2])
        stride = params.get('stride', kernel_size)
        
        # Handle scalar kernel_size/stride
        if isinstance(kernel_size, int):
            kernel_size = [kernel_size, kernel_size]
        if isinstance(stride, int):
            stride = [stride, stride]
        
        # Calculate output dimensions
        out_h = (in_h - kernel_size[0]) // stride[0] + 1
        out_w = (in_w - kernel_size[1]) // stride[1] + 1
        
        return [batch, channels, out_h, out_w]
    
    def _infer_adaptive_avgpool2d_shape(self, input_shape: List[int], node: Dict[str, Any]) -> List[int]:
        """Infer AdaptiveAvgPool2D output shape."""
        params = node.get('params', node.get('parameters', {}))
        
        batch = input_shape[0]
        channels = input_shape[1]
        
        # Get output size
        output_size = params.get('output_size', [1, 1])
        
        if isinstance(output_size, int):
            output_size = [output_size, output_size]
        
        return [batch, channels, output_size[0], output_size[1]]
    
    def _infer_reshape_shape(self, input_shape: List[int], node: Dict[str, Any]) -> List[int]:
        """Infer Reshape output shape."""
        params = node.get('params', node.get('parameters', {}))
        
        new_shape = params.get('shape', input_shape)
        
        # Handle -1 in shape (auto-infer dimension)
        if -1 in new_shape:
            total_elements = int(np.prod(input_shape))
            known_elements = int(np.prod([s for s in new_shape if s != -1]))
            auto_dim = total_elements // known_elements
            
            new_shape = [auto_dim if s == -1 else s for s in new_shape]
        
        return new_shape
    
    def validate_shapes(self, mir: Dict[str, Any]) -> bool:
        """
        Validate that all shapes are compatible throughout the graph.
        
        Returns:
            True if all shapes are valid, False otherwise
        """
        print("\n🔍 Validating shapes...")
        
        issues = []
        
        nodes = mir.get('nodes', mir.get('graph', {}).get('nodes', []))
        
        for node in nodes:
            if 'shape_info' not in node:
                continue
            
            shape_info = node['shape_info']
            input_shapes = shape_info.get('input_shapes', [])
            output_shape = shape_info.get('output_shape')
            
            # Check for None/invalid shapes
            if output_shape is None:
                issues.append(f"Node '{node['id']}' has no output shape")
            
            if any(s is None for s in input_shapes):
                issues.append(f"Node '{node['id']}' has unknown input shapes")
        
        if issues:
            print(f"❌ Shape validation failed:")
            for issue in issues:
                print(f"   - {issue}")
            return False
        else:
            print(f"✅ All shapes valid")
            return True


def apply_shape_inference(mir: Dict[str, Any]) -> Dict[str, Any]:
    """
    Convenience function to apply shape inference pass.
    
    Args:
        mir: MIR graph
    
    Returns:
        mir: MIR with shape inference applied
    """
    shape_pass = ShapeInferencePass()
    mir = shape_pass.run(mir)
    shape_pass.validate_shapes(mir)
    
    return mir


if __name__ == '__main__':
    # Example usage
    print("Shape Inference Pass")
    print("=" * 60)
    print("\nThis pass automatically:")
    print("  1. Propagates shapes through the MIR graph")
    print("  2. Inserts flatten operations before Linear layers")
    print("  3. Validates shape compatibility")
    print("  4. Annotates nodes with shape metadata")
    print("\nUsage:")
    print("  from passes.shape_inference import apply_shape_inference")
    print("  mir = apply_shape_inference(mir)")