import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional class LinearWithBandwidth(nn.Module): """ A custom linear layer that outputs both feature values and per-node bandwidth parameters. The bandwidth information allows reshaping the output for Conv2D operations by providing spatial structure information for each neuron. Args: in_features: Size of each input sample out_features: Number of output features (neurons) bandwidth_dim: Spatial dimension for bandwidth (height/width of the output grid) use_bias: If True, includes a bias term for the linear transformation """ def __init__( self, in_features: int, out_features: int, bandwidth_dim: int, use_bias: bool = True ): super().__init__() self.in_features = in_features self.out_features = out_features self.bandwidth_dim = bandwidth_dim # Standard linear layer parameters self.weight = nn.Parameter(torch.randn(out_features, in_features)) if use_bias: self.bias = nn.Parameter(torch.zeros(out_features)) else: self.register_parameter('bias', None) # Bandwidth parameters per neuron # Each neuron gets a bandwidth value that can be used for spatial structuring self.bandwidth = nn.Parameter(torch.ones(out_features)) self.reset_parameters() def reset_parameters(self): """Initialize parameters using Kaiming initialization.""" nn.init.kaiming_uniform_(self.weight, a=5**0.5) if self.bias is not None: nn.init.zeros_(self.bias) # Initialize bandwidth to ones nn.init.ones_(self.bandwidth) def forward(self, x: torch.Tensor): """ Forward pass. Args: x: Input tensor of shape (batch_size, in_features) Returns: Tuple of (features, bandwidth) where: - features: shape (batch_size, out_features) - bandwidth: shape (out_features,) - shared across batch """ features = F.linear(x, self.weight, self.bias) return features, self.bandwidth def forward_reshaped( self, x: torch.Tensor, target_height: Optional[int] = None, target_width: Optional[int] = None ): """ Forward pass that reshapes output for Conv2D compatibility. Args: x: Input tensor of shape (batch_size, in_features) target_height: Height of the output grid (optional) target_width: Width of the output grid (optional) Returns: Reshaped tensor of shape (batch_size, channels, height, width) where bandwidth information is used to structure the output """ features, bandwidth = self(x) batch_size = x.shape[0] # Determine spatial dimensions if target_height is not None and target_width is not None: assert target_height * target_width >= self.out_features, \ "Grid size must be at least out_features" h, w = target_height, target_width else: # Default to square grid h = w = int(torch.ceil(torch.sqrt(torch.tensor(self.out_features, dtype=torch.float)))) while h * w < self.out_features: w += 1 # Create output tensor channels = 1 # Single channel with bandwidth-modulated features output = torch.zeros(batch_size, channels, h, w, device=x.device, dtype=x.dtype) # Fill in features using bandwidth as a weighting/scaling factor for i in range(self.out_features): row = i // w col = i % w if row < h and col < w: # Apply bandwidth scaling to the feature output[:, 0, row, col] = features[:, i] * self.bandwidth[i] return output class ConvAfterLinear(nn.Module): """ Example module showing how to use LinearWithBandwidth followed by Conv2D. Args: in_features: Input feature size hidden_features: Hidden layer size (output of linear with bandwidth) grid_size: Spatial dimension for reshaping (height/width) out_channels: Number of Conv2D output channels """ def __init__( self, in_features: int, hidden_features: int, grid_size: int, out_channels: int ): super().__init__() self.linear_bw = LinearWithBandwidth( in_features=in_features, out_features=hidden_features, bandwidth_dim=grid_size ) self.conv = nn.Conv2d( in_channels=1, out_channels=out_channels, kernel_size=3, padding=1 ) self.grid_size = grid_size def forward(self, x: torch.Tensor): """ Args: x: Input tensor of shape (batch_size, in_features) Returns: Output tensor of shape (batch_size, out_channels, grid_size, grid_size) """ # Get features with bandwidth and reshape for conv reshaped = self.linear_bw.forward_reshaped( x, target_height=self.grid_size, target_width=self.grid_size ) # Apply conv2d out = self.conv(reshaped) return out if __name__ == "__main__": # Example usage batch_size = 4 in_features = 128 hidden_features = 64 grid_size = 8 out_channels = 16 print("=== LinearWithBandwidth Example ===") print(f"Input: ({batch_size}, {in_features})") print(f"Hidden: {hidden_features}, Grid: {grid_size}x{grid_size}") print() # Test basic linear with bandwidth layer = LinearWithBandwidth(in_features, hidden_features, grid_size) x = torch.randn(batch_size, in_features) features, bandwidth = layer(x) print(f"Features shape: {features.shape}") print(f"Bandwidth shape: {bandwidth.shape}") print(f"Bandwidth values: {bandwidth[:5]}") print() # Test reshaped output reshaped = layer.forward_reshaped(x, grid_size, grid_size) print(f"Reshaped for Conv2D: {reshaped.shape}") print() # Test full pipeline with Conv2D model = ConvAfterLinear(in_features, hidden_features, grid_size, out_channels) out = model(x) print(f"Final output shape: {out.shape}") print() # Verify gradients flow through bandwidth optimizer = torch.optim.Adam(model.parameters(), lr=0.001) loss = out.sum() loss.backward() print("Gradient check:") print(f" weight.grad: {model.linear_bw.weight.grad is not None}") print(f" bandwidth.grad: {model.linear_bw.bandwidth.grad is not None}") print(f" conv.weight.grad: {model.conv.weight.grad is not None}")