//  [python]
def __init__(self,
               z_dims=64,
               condition_bijector=False,
               bijector_type="iaf",
               encoder_type="dense",
               beta=1.0,
               beta_steps=0,
               min_kl=0,
               use_q_z_for_ais=False,
               dataset=None,
               prior_type="simple",
               affine_rank=1):
    self.train_size = dataset.train_size
    if dataset.name == "mnist":
      output_shape = [28, 28, 1]
    elif dataset.name == "cifar10":
      output_shape = [32, 32, 3]

    self._z_dims = z_dims
    self._beta = beta
    self._use_q_z_for_ais = use_q_z_for_ais
    if beta_steps > 0:
      frac = tf.cast(
          tf.train.get_or_create_global_step(), tf.float32) / tf.cast(beta_steps, tf.float32)
      frac = tf.minimum(frac, 1.0)
      self._beta = frac * beta
    else:
      self._beta = tf.constant(beta)
    self._min_kl = tf.cast(min_kl, tf.float32)

    if bijector_type == "affine":
      # TriL + shift
      num_outputs = (z_dims * (z_dims + 1)) // 2 + z_dims
    elif bijector_type == "affine_lr":
      num_outputs = z_dims * 2 + z_dims * affine_rank
    elif condition_bijector and bijector_type not in ["conv_shift_scale", "shift_scale"]:
      num_outputs = 3 * z_dims
    else:
      num_outputs = 2 * z_dims

    if encoder_type == "hier_conv":
      assert dataset.name == "cifar10"
      self._encoder = ConvHierEncoder("encoder")
      self._prior_posterior = ConvHierPriorPost("prior_post")
      self._decoder = lambda z: self._prior_posterior(z=z)[2]
      self._prior = lambda z=None, batch=None: self._prior_posterior(z=z, batch=batch)[:2]
      self._recog = lambda images, z=None: self._prior_posterior(images=images, z=z, encoder=self._encoder)[:2]
    else:
      if encoder_type == "dense":
        self._encoder = DenseEncoder(
            "encoder", num_outputs=num_outputs, activation=tf.nn.softplus)
        self._decoder = DenseDecoder(
            "decoder", output_shape=output_shape, activation=tf.nn.softplus)
      elif encoder_type == "conv":
        self._encoder = ConvEncoder("encoder", num_outputs=num_outputs)
        self._decoder = ConvDecoder("decoder", output_shape=output_shape)
        conv_z_shape = [4, 4, self._z_dims]
      elif encoder_type == "conv2":
        self._encoder = ConvEncoder2("encoder", num_outputs=num_outputs)
        self._decoder = ConvDecoder2("decoder", output_shape=output_shape)
        conv_z_shape = [4, 4, self._z_dims]
      elif encoder_type == "conv3":
        self._encoder = ConvEncoder3("encoder", num_outputs=num_outputs)
        self._decoder = ConvDecoder3("decoder", output_shape=output_shape)
        conv_z_shape = [8, 8, self._z_dims]
      elif encoder_type == "conv4":
        self._encoder = ConvEncoder4("encoder", num_outputs=num_outputs)
        self._decoder = ConvDecoder4("decoder", output_shape=output_shape)
        conv_z_shape = [8, 8, self._z_dims]

      if prior_type == "simple":
        self._prior = SimplePrior("prior", num_dims=self._z_dims)
      elif prior_type == "simple_3d":
        self._prior = Simple3DPrior("prior", shape=conv_z_shape)
      if bijector_type == "iaf":
        recog = DenseRecognitionIAF(
            "recog", encoder=self._encoder, condition_iaf=condition_bijector)
      elif bijector_type == "rnvp":
        recog = DenseRecognitionRNVP(
            "recog",
            encoder=self._encoder,
            condition_bijector=condition_bijector)
      elif bijector_type == "shift_scale":
        recog = DenseRecognition("recog", encoder=self._encoder)
      elif bijector_type == "conv_shift_scale":
        recog = ConvShiftScale("recog", encoder=self._encoder)
      elif bijector_type == "affine":
        recog = DenseRecognitionAffine("recog", encoder=self._encoder, z_dims=z_dims)
      elif bijector_type == "conv_iaf":
        recog = ConvIAF("recog", encoder=self._encoder, condition_iaf=condition_bijector)
      elif bijector_type == "affine_lr":
        recog = DenseRecognitionAffineLR("recog", encoder=self._encoder, z_dims=z_dims, rank=affine_rank)

      # Drop the bijector return.
      self._recog = lambda *args, **kwargs: recog(*args, **kwargs)[:2]

    if dataset.name == "mnist":
      self._noise = DenseMNISTNoise("noise", decoder=self._decoder)
    else:
      self._noise = DenseCIFAR10TNoise("noise", decoder=self._decoder)