class_docstr: Fits a PixelCNN in the VQVAE's latent space.
  
  Args:
      model (VQVAE): The AE model to sample from
      sampler_config (PixelCNNSamplerConfig): A PixelCNNSamplerConfig instance containing
          the main parameters of the sampler. If None, a pre-defined configuration is used.
          Default: None
  
  .. note::
  
      The method :class:`~pythae.samplers.MAFSampler.fit` must be called to fit the
      sampler before sampling.
class_docstr_tok: ["Fits","a","PixelCNN","in","the","VQVAE","'s","latent","space",".","Args",":","model","(","VQVAE",")",":","The","AE","model","to","sample","from","sampler_config","(","PixelCNNSamplerConfig",")",":","A","PixelCNNSamplerConfig","instance","containing","the","main","parameters","of","the","sampler",".","If","None",",","a","pre-defined","configuration","is","used",".","Default",":","None","..","note",":",":","The","method",":","class",":","`","~pythae.samplers.MAFSampler.fit","`","must","be","called","to","fit","the","sampler","before","sampling","."]
code: def __init__(
      self, model: VQVAE, sampler_config: PixelCNNSamplerConfig = None
  ):
      self.is_fitted = False
  
      if sampler_config is None:
          sampler_config = PixelCNNSamplerConfig()
  
      BaseSampler.__init__(
          self, model=model, sampler_config=sampler_config
      )
  
      # get size of codes
      x_dumb = torch.randn((2,) + model.model_config.input_dim).to(
          self.device
      )
      out_dumb = model({"data": x_dumb})
      quant_dumb = out_dumb.quantized_indices
      z_dumb = out_dumb.z
  
      self.needs_reshape = False
      if len(z_dumb.shape) == 2:
          self.needs_reshape = True
  
      pixelcnn_config = PixelCNNConfig(
          input_dim=quant_dumb.shape[1:],
          n_embeddings=model.model_config.num_embeddings,
          n_layers=sampler_config.n_layers,
          kernel_size=sampler_config.kernel_size,
      )
  
      self.pixelcnn_model = PixelCNN(model_config=pixelcnn_config).to(
          self.device
      )
  
file_path: benchmark_VAE/src/pythae/samplers/pixelcnn_sampler/pixelcnn_sampler.py
imports: import os
  import shutil
  from typing import Union
  import numpy
  import torch
  import torch.nn.functional
  from torch.utils.data import DataLoader, Dataset
  from ...data.datasets import collate_dataset_output
  from ...data.preprocessors import DataProcessor
  from ...models import VQVAE
  from ...models.normalizing_flows import PixelCNN, PixelCNNConfig
  from ...trainers import BaseTrainer, BaseTrainerConfig
  from ..base.base_sampler import BaseSampler
  from .pixelcnn_sampler_config import PixelCNNSamplerConfig
label_desc: Use image node_id 1 to create a new PixelCNNSampler object from inherited base classes: BaseSampler with example: obj = PixelCNNSampler(model, sampler_config)