| # Auto-extracted class source (static) | |
| class IcosahedralRRF(nn.Module): | |
| def __init__(self, input_dim, hidden_dim, output_dim, gnn_num_layers=2, gnn_z_dim=16, gnn_alpha_attn=1.0, gnn_dropout=0.1): | |
| super(IcosahedralRRF, self).__init__() | |
| # 12 nodos gauge | |
| self.nodes = nn.ModuleList([ | |
| SavantRRF_Gauge(input_dim, hidden_dim, output_dim) for _ in range(12) | |
| ]) | |
| # Núcleo ético | |
| # The input to ethical_core is the concatenation of the outputs of the 12 gauge nodes. | |
| # Each gauge node outputs a tensor of shape [batch_size, output_dim]. | |
| # Concatenating these along dim=1 results in a shape [batch_size, 12 * output_dim]. | |
| self.ethical_core = nn.Linear(12 * output_dim, output_dim) | |
| # Subconsciente (dodecaedro) using GNNDiracRRF | |
| # The input dimension (in_dim) for the GNN should match the feature dimension of its input nodes. | |
| # There's ambiguity in the original code about what the GNN's nodes and features are. | |
| # Interpretation 1 (based on original code passing 'regulated'): GNN operates on 'batch_size' nodes, with 'output_dim' features. in_dim = output_dim. | |
| # Interpretation 2 (more conventional for graph on icosahedron/dodecahedron): GNN operates on 12 or 20 nodes, with features derived from gauge outputs. | |
| # Let's assume interpretation 2, where the GNN operates on the 12 gauge nodes. | |
| # The features for each of these 12 nodes would be the output of the corresponding gauge node, shape [batch_size, output_dim]. | |
| # For a GNN layer expecting [num_nodes, in_channels], the input should be [12, output_dim] per batch item. | |
| # This means the GNN's in_dim should be output_dim. This matches the current GNN init below. | |
| # The GNN's out_dim should match the desired output feature dimension per node (e.g., output_dim). | |
| # The number of nodes for the GNN is 12 (for icosahedral). | |
| # Let's define the memory_map GNN assuming it operates on the 12 gauge nodes. | |
| # The input features to the GNN will be the outputs of the 12 gauge nodes. | |
| # Each gauge node outputs a tensor of shape [batch_size, output_dim]. | |
| # We will treat output_dim as the feature dimension for the GNN nodes (the 12 gauge nodes). | |
| # So, in_dim for GNN = output_dim. | |
| # The GNN will output features for each of the 12 nodes. Let's assume out_dim for GNN is also output_dim. | |
| self.memory_map = GNNDiracRRF(in_dim=output_dim, # Feature dimension for GNN nodes (output_dim of gauge nodes) | |
| hidden_dim=hidden_dim, | |
| out_dim=output_dim, # Output feature dimension per GNN node | |
| num_layers=gnn_num_layers, | |
| z_dim=gnn_z_dim, | |
| alpha_attn=gnn_alpha_attn, | |
| dropout=gnn_dropout) | |
| def forward(self, x, edge_index=None, z=None): | |
| # x is the input to the gauge nodes, shape [batch_size, input_dim, sequence_length] | |
| outputs = [node(x) for node in self.nodes] | |
| # outputs is a list of 12 tensors, each [batch_size, output_dim] | |
| # Concatenate outputs for the ethical core | |
| concat = torch.cat(outputs, dim=1) # [batch_size, 12 * output_dim] | |
| regulated = torch.sigmoid(self.ethical_core(concat)) # [batch_size, output_dim] | |
| # GNN operation on the 12 gauge nodes | |
| if edge_index is not None and z is not None: | |
| # Prepare input for the GNN: Features for the 12 nodes (the gauge node outputs). | |
| # Stack the outputs to get [batch_size, 12, output_dim] | |
| stacked_outputs = torch.stack(outputs, dim=1) # [batch_size, 12, output_dim] | |
| # Reshape for GNN input: [num_nodes, in_channels] = [12, output_dim] per batch item. | |
| # Need to process batch items. Simplest is to iterate. | |
| # A more efficient way is to use torch_geometric.data.Batch | |
| gnn_outputs_list = [] | |
| for i in range(stacked_outputs.size(0)): | |
| # GNN input features for this batch item: [12, output_dim] | |
| gnn_input_features_i = stacked_outputs[i] | |
| # Ensure edge_index and z are on the correct device | |
| edge_index_i = edge_index.to(x.device) | |
| z_i = z.to(x.device) | |
| # GNN forward pass for one batch item | |
| gnn_output_i = self.memory_map(gnn_input_features_i, edge_index_i, z_i) # [12, output_dim] | |
| gnn_outputs_list.append(gnn_output_i) | |
| # Stack GNN outputs back into a batch tensor: [batch_size, 12, output_dim] | |
| gnn_outputs_stacked = torch.stack(gnn_outputs_list, dim=0) | |
| # Now, how to combine the GNN output [batch_size, 12, output_dim] with the 'regulated' output [batch_size, output_dim]? | |
| # The original model returned just 'regulated'. | |
| # A simple approach is to maybe combine them, e.g., add, concatenate, or use the GNN output as a modulation. | |
| # Let's stick to returning the aggregated GNN output as the final output when GNN is used. | |
| # This changes the model's behavior compared to the original. | |
| # Alternative: The GNN output modulates the 'regulated' output. | |
| # E.g., regulated * sigmoid(aggregated_gnn_output) or similar. | |
| # Let's stick to returning the aggregated GNN output when edge_index and z are provided, | |
| # and the original 'regulated' output otherwise. This seems the most direct path based on the conditional in the original forward. | |
| # Aggregate the 12 nodes' outputs from the GNN | |
| aggregated_gnn_output = gnn_outputs_stacked.mean(dim=1) # [batch_size, output_dim] | |
| return aggregated_gnn_output # [batch_size, output_dim] | |
| else: | |
| # If edge_index and z are not provided, return the output of the ethical core as before. | |
| return regulated |