Project: s41467-025-56424-6 | Function: negative_binomial_loss (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:** Implementation of negative binomial loss function for reconstructing RNA and ADT count data. Computes the negative log-likelihood under negative binomial distribution with mean and dispersion parameters, scaled by size factors.

**KEY TERMS:** {
  "mean": "reconstructed mean expression",
  "disp": "dispersion parameter",
  "scale_factor": "size factor normalization",
  "x": "observed count data",
  "eps": "numerical stability constant"
}

**CONTEXT CODE:**             [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 = []
        for i in range(self.omics_num):
            hiddens.append(self.encoders[i](inputs[i]))
        return hiddens

    def decode(self, hiddens):
        means, disps = [], []
        for i in range(self.omics_num):
            mean, disp = self.decoders[i](hiddens[i])
            means.append(mean)
            disps.append(disp)
        return means, disps

    def decode_q(self, hidden):
        means, disps = [], []
        for i in range(self.omics_num):
            mean, disp = self.decoders_q[i](hidden)
            means.append(mean)
            disps.append(disp)
        return means, disps


def negative_binomial_loss(mean, disp, scale_factor, x):
    # MASKED: Implementation of negative binomial loss function for reconstructing RNA and ADT count data. Computes the negative log-likelihood under negative binomial distribution with mean and dispersion parameters, scaled by size factors.
    raise NotImplementedError("This function is masked for testing purposes.")


def poisson_loss(mean, scale_factor, x):
    eps = 1e-12
    mean = mean * scale_factor

    poisson = mean - x * torch.log(mean + eps) + torch.lgamma(x + 1)

    return poisson


def reconstruction_loss(mean, disp, scale_factor, x, data_type, reduction="mean"):
    if data_type == "RNA" or data_type == "ADT":
        loss = negative_binomial_loss(mean, disp, scale_factor, x)
    elif data_type == "ATAC":
        loss = poisson_loss(mean, scale_factor, x)
    if reduction == "mean":
        loss = torch.mean(loss)
    elif reduction == "sum":
        loss = torch.sum(loss)
    else:
        raise NotImplementedError
    return loss


Complete the function implementation replacing "raise NotImplementedError".