""" ConvMLP Utilities ================= Vedic-inspired optimized utility functions. """ import numpy as np from typing import Tuple, Union def im2col(x: np.ndarray, kernel_size: Tuple[int, int], strides: Tuple[int, int] = (1, 1)) -> np.ndarray: """ Image to Column conversion (Vedic Block Multiplication). Inspired by Urdhva-Tiryagbyham (vertically and crosswise). Converts 2D convolution to efficient matrix multiplication. Parameters ---------- x : np.ndarray Input tensor of shape (batch, channels, height, width) kernel_size : tuple (kernel_height, kernel_width) strides : tuple (stride_h, stride_w) Returns ------- col : np.ndarray Unfolded patches of shape (batch, channels * kh * kw, out_h * out_w) Example ------- >>> patches = im2col(x, kernel_size=(3, 3), strides=(1, 1)) >>> # patches[:, :, 0] = first patch flattened """ batch, channels, h, w = x.shape kh, kw = kernel_size sh, sw = strides out_h = (h - kh) // sh + 1 out_w = (w - kw) // sw + 1 # Efficient indexing using strided views col = np.zeros((batch, channels, kh, kw, out_h, out_w), dtype=x.dtype) for i in range(out_h): for j in range(out_w): h_start = i * sh w_start = j * sw col[:, :, :, :, i, j] = x[:, :, h_start:h_start+kh, w_start:w_start+kw] # Reshape to (batch, channels * kh * kw, out_h * out_w) col = col.transpose(0, 4, 5, 1, 2, 3).reshape(batch, out_h * out_w, channels * kh * kw) return col.transpose(0, 2, 1) def col2im(col: np.ndarray, x_shape: Tuple[int, ...], kernel_size: Tuple[int, int], strides: Tuple[int, int] = (1, 1)) -> np.ndarray: """ Column to Image conversion (Reverse of im2col). Parameters ---------- col : np.ndarray Column data of shape (batch, channels * kh * kw, out_h * out_w) x_shape : tuple Original input shape (batch, channels, height, width) kernel_size : tuple (kernel_height, kernel_width) strides : tuple (stride_h, stride_w) Returns ------- x : np.ndarray Reconstructed image of shape x_shape """ batch, channels, h, w = x_shape kh, kw = kernel_size sh, sw = strides out_h = (h - kh) // sh + 1 out_w = (w - kw) // sw + 1 # Reshape col back to (batch, out_h, out_w, channels, kh, kw) col = col.transpose(0, 2, 1).reshape(batch, out_h, out_w, channels, kh, kw) # Initialize output with zeros x = np.zeros(x_shape, dtype=col.dtype) # Accumulate values for i in range(out_h): for j in range(out_w): h_start = i * sh w_start = j * sw x[:, :, h_start:h_start+kh, w_start:w_start+kw] += col[:, i, j] return x def calculate_receptive_field(layers: list, start_rf: int = 1) -> list: """ Calculate receptive field at each layer. Based on the "coconut sellers" concept - each layer adds to the receptive field. Parameters ---------- layers : list List of layer objects with kernel_size and strides attributes start_rf : int Initial receptive field size Returns ------- receptive_fields : list Receptive field size at each layer """ receptive_fields = [start_rf] current_rf = start_rf for layer in layers: if hasattr(layer, 'kernel_size') and hasattr(layer, 'strides'): kh, kw = layer.kernel_size sh, sw = layer.strides if isinstance(layer.strides, tuple) else (layer.strides, layer.strides) current_rf = current_rf + (kh - 1) * np.prod([l.strides if isinstance(l, tuple) else l for l in receptive_fields]) receptive_fields.append(current_rf) elif hasattr(layer, 'pool_size'): ph, pw = layer.pool_size if isinstance(layer.pool_size, tuple) else (layer.pool_size, layer.pool_size) current_rf = current_rf * ph receptive_fields.append(current_rf) return receptive_fields def next_power_of_2(n: int) -> int: """Return the next power of 2 >= n.""" return 2 ** int(np.ceil(np.log2(n))) def calculate_flops(conv_layers: list, input_shape: Tuple[int, ...]) -> int: """ Calculate total FLOPs for a convolutional network. Inspired by Bhaskara's wheel method - counting all operations. """ total_flops = 0 current_shape = input_shape for layer in conv_layers: if isinstance(layer, Conv2D): kh, kw = layer.kernel_size out_h, out_w = layer._compute_output_shape(current_shape)[-2:] # FLOPs = output_elements * kernel_elements * input_channels flops = out_h * out_w * kh * kw * current_shape[-3] * layer.filters total_flops += flops current_shape = layer._compute_output_shape(current_shape) elif isinstance(layer, (MaxPool2D, AvgPool2D)): ph, pw = layer.pool_size if isinstance(layer.pool_size, tuple) else (layer.pool_size, layer.pool_size) current_shape = (current_shape[0], current_shape[1], current_shape[2] // ph, current_shape[3] // pw) return total_flops class NikhilamQuantizer: """ Nikhilam Quantization for fast inference. Based on the Vedic "deficiency from base" method. Quantizes weights to low-precision for faster computation. """ def __init__(self, num_bits: int = 8): self.num_bits = num_bits self.scale = None self.zero_point = None def fit(self, weights: np.ndarray): """Calculate quantization parameters.""" self.scale = (weights.max() - weights.min()) / (2 ** self.num_bits - 1) self.zero_point = weights.min() return self def quantize(self, weights: np.ndarray) -> np.ndarray: """Quantize weights to low precision.""" if self.scale is None: self.fit(weights) return np.round((weights - self.zero_point) / self.scale).astype(np.int8) def dequantize(self, weights: np.ndarray) -> np.ndarray: """Restore weights from low precision.""" return weights.astype(np.float32) * self.scale + self.zero_point def validate_input_shape(x: np.ndarray, expected_channels: int = None) -> bool: """Validate input tensor shape.""" if len(x.shape) == 2: # Flattened - need to reshape return False elif len(x.shape) == 3: # (channels, height, width) - needs batch dimension return False elif len(x.shape) == 4: # (batch, channels, height, width) - valid if expected_channels and x.shape[1] != expected_channels: raise ValueError(f"Expected {expected_channels} channels, got {x.shape[1]}") return True else: raise ValueError(f"Invalid input shape: {x.shape}")