def forward(self,
                x: torch.Tensor,
                data_format='channel_first') -> torch.Tensor:
        """Sparse layer norm forward function with 2D data.

        Args:
            x (torch.Tensor): The input tensor.
            data_format (str): The format of the input tensor. If
                ``"channel_first"``, the shape of the input tensor should be
                (B, C, H, W). If ``"channel_last"``, the shape of the input
                tensor should be (B, H, W, C). Defaults to "channel_first".
        """
        assert x.dim() == 4, (
            f'LayerNorm2d only supports inputs with shape '
            f'(N, C, H, W), but got tensor with shape {x.shape}')
        if data_format == 'channel_last':
            index = SparseHelper._get_active_map_or_index(
                H=x.shape[1], returning_active_map=False)

            # select the features on non-masked positions to form flatten
            # features with shape (n, c)
            x_flattened = x[index]
            # use LayerNorm to normalize this flatten feature (n, c)
            x_flattened = super().forward(x_flattened)

            # generate output
            x = torch.zeros_like(x, dtype=x_flattened.dtype)
            x[index] = x_flattened
        elif data_format == 'channel_first':
            index = SparseHelper._get_active_map_or_index(
                H=x.shape[2], returning_active_map=False)
            x_permuted = x.permute(0, 2, 3, 1)

            # select the features on non-masked positions to form flatten
            # features with shape (n, c)
            x_flattened = x_permuted[index]
            # use LayerNorm to normalize this flatten feature (n, c)
            x_flattened = super().forward(x_flattened)

            # generate output
            x = torch.zeros_like(x_permuted, dtype=x_flattened.dtype)
            x[index] = x_flattened
            x = x.permute(0, 3, 1, 2).contiguous()
        else:
            raise NotImplementedError
        return x
Sparse layer norm forward function with 2D data.

        Args:
            x (torch.Tensor): The input tensor.
            data_format (str): The format of the input tensor. If
                ``"channel_first"``, the shape of the input tensor should be
                (B, C, H, W). If ``"channel_last"``, the shape of the input
                tensor should be (B, H, W, C). Defaults to "channel_first".
        
mmpretrain/models/utils/sparse_modules.py