Project: s41467-024-46946-w | Function: forward (python)

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:** Implements the core DMNet deep neural network architecture for double-secure optical decryption, processing QR code security keys and speckle ciphertexts through complex-valued neural layers to recover plaintext images

**KEY TERMS:** {
  "x1": "QR code security key input",
  "x2": "speckle pattern ciphertext input",
  "fc_qr2spk": "complex linear layer f₁ transforming QR to speckle domain",
  "unet": "complex U-Net f_diff for speckle subtraction processing",
  "fc_out": "complex linear layer f_out for final plaintext reconstruction",
  "dropout1": "regularization for QR processing branch",
  "norm1": "normalization for synthetic QR speckle patterns"
}

**CONTEXT CODE:** import torch
import torch.nn as nn
import torch.nn.functional as F
from complexLayers import ComplexBatchNorm2d, ComplexConv2d, ComplexLinear, ComplexDropout
from complexFunctions import complex_relu, complex_max_pool2d

# Simple SwitchNorm2d implementation for testing
class SwitchNorm2d(nn.Module):
    def __init__(self, num_features):
        super(SwitchNorm2d, self).__init__()
        self.num_features = num_features
        self.bn = nn.BatchNorm2d(num_features)

    def forward(self, x):
        return self.bn(x)

class DMNet(nn.Module):
    '''
    optical propagation:
    image1 + image2 = speckles

    image1 is QR code
    image2 are from fMNIST.

        (input1)x1:qr       (input2)x2:full-speckle
            |                       |
    f1(qr):qr_speckle               |
            |                       |
            ------------|------------            
                        |
            x_diff = f_diff(f1(input1),input2): (full-speckle-qr_speckle=image_speckle)
                        |
                y = f_out(x_diff): image generated from image_speckle

    f1:     Complex_fc
    f_diff: Complex_Unet
    f_out:  Comlex_fc
    '''
    def __init__(self, in1_dim, in2_dim, out_dim, dropout):
        super(DMNet, self).__init__()
        self.in1_dim = in1_dim
        self.in2_dim = in2_dim
        self.out_dim = out_dim
        self.p = dropout
        self.fc_qr2spk = ComplexLinear(self.in1_dim*self.in1_dim, self.in2_dim*self.in2_dim) #f1
        self.fc_out = ComplexLinear(self.in2_dim*self.in2_dim, self.out_dim*self.out_dim) #f2
        self.norm1 = SwitchNorm2d(1)#   nn.InstanceNorm2d(1, affine=True)#nn.InstanceNorm2d(self.in2_dim*self.in2_dim)
        self.norm2 = SwitchNorm2d(1)#nn.InstanceNorm2d(1, affine=True)#nn.LayerNorm(self.in2_dim*self.in2_dim)
        self.norm3 = nn.LayerNorm(self.out_dim*self.out_dim)
        self.unet = Complex_Unet(dropout=self.p, inputCH=2)
        self.dropout1 = ComplexDropout(p=self.p)
        self.dropout2 = ComplexDropout(p=self.p)
        self.dp = nn.Dropout2d(p=self.p) #fine-tune : p=0.8
        self.sigmoid = nn.Sigmoid()
        self.bn = SwitchNorm2d(1)#nn.BatchNorm2d(1)
        #self.res = resnext101() #f_diff

        
    def forward(self, x1, x2):
        # MASKED: Implements the core DMNet deep neural network architecture for double-secure optical decryption, processing QR code security keys and speckle ciphertexts through complex-valued neural layers to recover plaintext images
        raise NotImplementedError("This function is masked for testing purposes.")

class Complex_Unet(torch.nn.Module):
    def __init__(self, dropout=0.5, GAN_type=1, inputCH=2):
        class Real_to_Complex(torch.nn.Module): # Transform Real number to Complex number
            def forward(self, input): 
                if GAN_type == 1: return input, torch.zeros(input.shape, dtype=input.dtype, device=input.device)
                if GAN_type == 0: return input
        
        class Complex_to_Real(torch.nn.Module): # Transform Complex number to Real number
            def forward(self, input):
                if GAN_type == 1: return torch.sqrt(input[0].square() + input[1].square())
                if GAN_type == 0: return input

        class Complex_conv2d(torch.nn.Module): # Complex Conv2d, perform Conv2d independently on real and imaginary part.
            def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True):
                super(Complex_conv2d, self).__init__()
                if GAN_type == 1: self.conv_r, self.conv_i = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias), torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
                if GAN_type == 0: self.conv_r              = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
            def forward(self, input): 
                if GAN_type == 1: return self.conv_r(input[0])-self.conv_i(input[1]), self.conv_r(input[0])+self.conv_i(input[1])
                if GAN_type == 0: return self.conv_r(input)

        class Complex_ConvTranspose2d(torch.nn.Module):
            def __init__(self,in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros'):
                super(Complex_ConvTranspose2d, self).__init__()
                if GAN_type == 1: self.conv_tran_r, self.conv_tran_i = torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias, dilation, padding_mode), torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias, dilation, padding_mode)
                if GAN_type == 0: self.conv_tran_r                   = torch.nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding, groups, bias, dilation, padding_mode)
            def forward(self, input):
                if GAN_type == 1: return self.conv_tran_r(input[0])-self.conv_tran_i(input[1]), self.conv_tran_r(input[0])+self.conv_tran_i(input[1])
                if GAN_type == 0: return self.conv_tran_r(input)

        class Complex_ReLU(torch.nn.Module): # Complex ReLU, perform ReLU independently on real and imaginary part.
            def __init__(self, inplace=False):
                super(Complex_ReLU, self).__init__()
                if GAN_type == 1: self.ReLU_r, self.ReLU_i = torch.nn.ReLU(inplace),torch.nn.ReLU(inplace)
                if GAN_type == 0: self.ReLU_r              = torch.nn.ReLU(inplace)
            def forward(self, input):
                if GAN_type == 1: return self.ReLU_r(input[0]), self.ReLU_i(input[1])
                if GAN_type == 0: return self.ReLU_r(input)

        class Complex_batch_norm2d(torch.nn.Module): # Complex BatchNorm, perform BatchNorm independently on real and imaginary part.
            def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True):
                super(Complex_batch_norm2d, self).__init__()
                if GAN_type == 1: self.bn_r, self.bn_i = torch.nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats),torch.nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats)
                if GAN_type == 0: self.bn_r            = torch.nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats)
            def forward(self, input): 
                if GAN_type == 1: return self.bn_r(input[0]), self.bn_i(input[1])
                if GAN_type == 0: return self.bn_r(input)

        class Complex_dropout2d(torch.nn.Module): # Complex Dropout, perform Dropout independently on real and imaginary part.
            def __init__(self,p=dropout, inplace=False):
                super(Complex_dropout2d,self).__init__()


Complete the function implementation replacing "raise NotImplementedError".