Project: 048_Topological_Blindspots_Understanding_and_Extending | Function: SubComplexLowConv.forward (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Core SMCN (Scalable Multi-Cellular Networks) forward pass implementing the Subcomplex Layer (SCL) update that processes bags of augmented Hasse graphs to address HOMP expressivity limitations through multi-cellular cochain spaces

**KEY TERMS:** {
  "SubComplexLowConv": "SMCN subcomplex layer",
  "edge_index_low_adjacency_subcomplex": "augmented Hasse graph adjacency",
  "bridge_cell_values": "bridge cell features",
  "conv": "expressive graph layer (GIN/GINE)",
  "edge_attributes": "edge features from bridge cells"
}

**CONTEXT CODE:** from typing import Tuple, Optional
from data.complex_data import ComplexData
from models.layers.homp import CustomGINConv


class SubComplexLowConv(nn.Module):
    """
    updates the x_i_j feature using the bag of graphs induced by the low rank cells
    """

    def __init__(
        self,
        low_rank: int,
        high_rank: int,
        input_channels: int,
        output_channels: int,
        hidden_channels: int = 16,
        number_of_mlp_layers: int = 2,
        conv_type: str = "custom_gin",
        use_batch_norm: bool = True,
        add_residual: bool = False,
        use_second_conv: bool = False,
        second_conv_type: Optional[str] = None,
        edge_dim: Optional[int] = None,
        dropout: float = 0.0,
        train_eps: bool = True,
        activation: str = "relu",
    ):
        super().__init__()
        self.input_channels = input_channels
        self.low_rank = low_rank
        self.high_rank = high_rank
        self.conv_type = conv_type
        self.edge_dim = edge_dim
        self.conv = self._get_conv(
            input_channels=input_channels,
            hidden_channels=hidden_channels,
            use_batch_norm=use_batch_norm,
            conv_type=conv_type,
            edge_dim=edge_dim,
            train_eps=train_eps,
            activation=activation,
        )
        self.use_second_conv = use_second_conv
        if use_second_conv:
            self.second_conv = self._get_second_conv(second_conv_type)

        self.bridge_rank = 1  # currently we only support low brifge rank = 1
        self.add_residual = add_residual

    def forward(self, data: Data) -> torch.Tensor:
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def get_edge_attributes(self, data: ComplexData) -> torch.Tensor:
        # computes the edge attributes based on bridge cell features
        bridge_cell_values = getattr(data, f"x_{self.bridge_rank}")
        indices = getattr(
            data,
            f"bridge_index_low_adjacency_subcomplex_{self.low_rank}_{self.high_rank}",
        )
        edge_attributes = bridge_cell_values[indices, :]
        return edge_attributes

    @staticmethod
    def _get_conv(
        input_channels: int,
        hidden_channels: int,
        use_batch_norm: bool,
        conv_type: str,
        edge_dim: Optional[int] = None,
        train_eps: bool = True,
        activation: str = "relu",
    ):

        mlp = get_mlp(
            input_dim=input_channels,
            hidden_dim=hidden_channels,
            batch_norm=use_batch_norm,
            activation=activation,
        )

        if "gin" == conv_type:
            return GINConv(mlp, train_eps=train_eps)
        elif "gine" == conv_type:
            return GINEConv(mlp, train_eps=train_eps, edge_dim=edge_dim)
        elif "custom_gin" == conv_type:
            return CustomGINConv(mlp, train_eps=train_eps, edge_dim=edge_dim)
        else:
            raise NotImplementedError("Only GIN convs are currently implemented...")

    def _get_second_conv(self, conv_type: str):
        if "repeat" == conv_type:
            return self.conv
        elif "id" == conv_type:
            return GINConv(nn.Identity(), train_eps=True)
        else:
            raise NotImplementedError("not implemented...")


class SubComplexIncidenceConv(nn.Module):
    """
    updates the x_i_j feature using the bag of graphs induced by the low rank cells
    """

    def __init__(
        self,
        low_rank: int,
        high_rank: int,
        input_channels: int,
        output_channels: int,
        hidden_channels: int = 16,
        conv_type: str = "gin",
        use_batch_norm: bool = True,
        add_residual: bool = False,
        use_second_conv: bool = False,
        second_conv_type: Optional[str] = None,
        number_of_mlp_layers: int = 2,
        dropout: float = 0.0,
        train_eps: bool = True,
        activation: str = "relu",
    ):
        super().__init__()
        self.input_channels = input_channels
        self.low_rank = low_rank
        self.high_rank = high_rank
        self.conv = self._get_conv(
            input_channels=input_channels,
            hidden_channels=hidden_channels,
            use_batch_norm=use_batch_norm,
            conv_type=conv_type,
            train_eps=train_eps,
            activation=activation,
        )
        self.conv_type = conv_type
        self.add_residual = add_residual

        self.use_second_conv = use_second_conv
        if use_second_conv:
            self.second_conv = self._get_second_conv(second_conv_type)

    def forward(self, data: Data) -> torch.Tensor:
        # Get data of input and output ranks
        x = getattr(data, f"x_{self.low_rank}_{self.high_rank}")
        edge_index = getattr(
            data, f"edge_index_incidence_subcomplex_{self.low_rank}_{self.high_rank}"
        )
        h = self.conv(x=x, edge_index=edge_index)
        x = x + h if self.add_residual else h
        x = self.second_conv(x=x, edge_index=edge_index) if self.use_second_conv else x
        return x


Complete the function implementation replacing "raise NotImplementedError".