Change all occurrences to _MODULEKEYS

Author: LouiseDck

scgen/_scgenvae.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch
 from scvi import REGISTRY_KEYS
+from scvi.module._constants import _MODULEKEYS
 from scvi.module.base import BaseModuleClass, LossOutput, auto_move_data
 from scvi.nn import Encoder
 from torch.distributions import Normal
@@ -80,15 +81,15 @@ def __init__(
 
     def _get_inference_input(self, tensors):
         x = tensors[REGISTRY_KEYS.X_KEY]
-        input_dict = dict(
-            x=x,
-        )
+        input_dict = {
+            REGISTRY_KEYS.X_KEY:x,
+        }
         return input_dict
 
     def _get_generative_input(self, tensors, inference_outputs):
-        z = inference_outputs["z"]
+        z = inference_outputs[_MODULEKEYS.Z_KEY]
         input_dict = {
-            "z": z,
+            _MODULEKEYS.Z_KEY: z,
         }
         return input_dict
 
@@ -101,15 +102,15 @@ def inference(self, x):
         """
         qz_m, qz_v, z = self.z_encoder(x)
 
-        outputs = dict(z=z, qzm=qz_m, qzv=qz_v)
+        outputs = {_MODULEKEYS.Z_KEY:z, _MODULEKEYS.QZM_KEY:qz_m, _MODULEKEYS.QZV_KEY:qz_v}
         return outputs
 
     @auto_move_data
     def generative(self, z):
         """Runs the generative model."""
         px = self.decoder(z)
 
-        return dict(px=px)
+        return {_MODULEKEYS.PX_KEY:px}
 
     def loss(
         self,
@@ -118,9 +119,9 @@ def loss(
         generative_outputs,
     ):
         x = tensors[REGISTRY_KEYS.X_KEY]
-        qz_m = inference_outputs["qzm"]
-        qz_v = inference_outputs["qzv"]
-        p = generative_outputs["px"]
+        qz_m = inference_outputs[_MODULEKEYS.QZ_M_KEY]
+        qz_v = inference_outputs[_MODULEKEYS.QZ_V_KEY]
+        p = generative_outputs[_MODULEKEYS.PX_KEY]
 
         kld = kl(
             Normal(qz_m, torch.sqrt(qz_v)),
@@ -164,7 +165,7 @@ def sample(
             inference_kwargs=inference_kwargs,
             compute_loss=False,
         )
-        px = Normal(generative_outputs["px"], 1).sample()
+        px = Normal(generative_outputs[_MODULEKEYS.PX_KEY], 1).sample()
         return px.cpu().numpy()
 
     def get_reconstruction_loss(self, x, px) -> torch.Tensor: