Project: s41467-025-56424-6 | 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:** Core vector quantization algorithm that assigns cells to discrete codebook entries and computes quantization loss. Takes cell embeddings as input, performs cosine similarity-based assignment to codebook entries, applies quantization, and updates entry usage statistics with rebalancing mechanisms for small/large entries.

**KEY TERMS:** {
  "assignment_indices": "cell-to-metacell assignments",
  "sim": "cosine similarity matrix",
  "entry.weight": "codebook entries (metacells)",
  "entry_prob": "entry usage statistics",
  "loss": "codebook loss L_C"
}

**CONTEXT CODE:** import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


class Quantizer(nn.Module):
    def __init__(self, entry_num, entry_dim):
        super().__init__()

        self.entry_num = entry_num
        self.entry_dim = entry_dim
        self.decay = 0.9
        self.entry = nn.Embedding(self.entry_num, self.entry_dim)
        self.register_buffer("entry_prob", torch.zeros(self.entry_num))

    def init_codebook(self, z, method):
        if method == "Random":
            self.entry.weight.data.uniform_(-1.0 / self.entry_num, 1.0 / self.entry_num)
        elif method == "Kmeans":
            import faiss

            d = z.shape[1]
            kmeans = faiss.Kmeans(d, self.entry_num, spherical=True, gpu=True)
            kmeans.train(z)
            D, I = kmeans.index.search(z, 1)
            assignments = I.reshape(-1)
            centers = np.zeros((self.entry_num, d))
            for i in range(self.entry_num):
                centers[i] = z[assignments == i].mean(axis=0)
            self.entry.weight.data.copy_(torch.from_numpy(centers))
        elif method == "Geometric":
            from geosketch import gs

            sketch_index = gs(z, self.entry_num, replace=False)
            self.entry.weight.data.copy_(torch.from_numpy(z[sketch_index]))

    def forward(self, e, return_assignment):
        # MASKED: Core vector quantization algorithm that assigns cells to discrete codebook entries and computes quantization loss. Takes cell embeddings as input, performs cosine similarity-based assignment to codebook entries, applies quantization, and updates entry usage statistics with rebalancing mechanisms for small/large entries.
        raise NotImplementedError("This function is masked for testing purposes.")


class Encoder(nn.Module):
    def __init__(self, input_dim, entry_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Linear(128, entry_dim),
        )

    def forward(self, input):
        return self.encoder(input)


class Decoder(nn.Module):
    def __init__(self, entry_dim, output_dim):
        super().__init__()
        self.docoder = nn.Sequential(
            nn.Linear(entry_dim, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Linear(128, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
        )
        self.decoder_mean = nn.Linear(512, output_dim)
        self.decoder_disp = nn.Sequential(
            nn.Linear(512, output_dim),
            nn.Softplus(),
        )

    def forward(self, input):
        decode = self.docoder(input)
        mean = torch.clamp(torch.exp(self.decoder_mean(decode)), 1e-5, 1e6)
        disp = torch.clamp(self.decoder_disp(decode), 1e-4, 1e4)
        return mean, disp


class Decoder_ATAC(nn.Module):
    def __init__(self, entry_dim, output_dim):
        super().__init__()
        self.docoder = nn.Sequential(
            nn.BatchNorm1d(entry_dim),
            nn.Linear(entry_dim, output_dim),
        )

    def forward(self, input):
        mean = torch.clamp(torch.exp(self.docoder(input)), 1e-5, 1e6)
        return mean, None


def get_decoder(data_type, entry_dim, output_dim):
    if data_type == "RNA" or data_type == "ADT":
        return Decoder(entry_dim, output_dim)
    elif data_type == "ATAC":
        return Decoder_ATAC(entry_dim, output_dim)


class MetaQ(nn.Module):
    def __init__(self, input_dims, data_types, entry_num, entry_dim=32):
        super(MetaQ, self).__init__()
        self.omics_num = len(input_dims)
        self.encoders = nn.ModuleList(
            [Encoder(input_dim, entry_dim) for input_dim in input_dims]
        )
        self.quantizer = Quantizer(entry_num, entry_dim * self.omics_num)
        self.decoders = nn.ModuleList(
            [
                get_decoder(data_type, entry_dim, input_dim)
                for input_dim, data_type in zip(input_dims, data_types)
            ]
        )
        self.decoders_q = nn.ModuleList(
            [
                get_decoder(data_type, entry_dim * self.omics_num, input_dim)
                for input_dim, data_type in zip(input_dims, data_types)
            ]
        )

    def copy_decoder_q(self):
        if self.omics_num == 1:
            self.decoders_q[0].load_state_dict(self.decoders[0].state_dict())

    def quantize(self, hiddens, return_assignment=False):
        if self.omics_num > 1:
            hidden = torch.cat(hiddens, dim=1)
        else:
            hidden = hiddens[0]
        return self.quantizer(hidden, return_assignment)

    def forward(self, inputs):
        hiddens = []


Complete the function implementation replacing "raise NotImplementedError".