class_docstr: Defines the ImageVAE model.
class_docstr_tok: ["Defines","the","ImageVAE","model","."]
code: def vae_internal(self, features, hparams, train):
      # inputs and targets should both be images with dims [batch, 64, 64, 1]
      inputs, targets = features["inputs"], features["targets"]
      inputs = tf.reshape(inputs, [-1, 64, 64, 1])
      targets = tf.reshape(targets, [-1, 64, 64, 1])
  
      clss = features["cls"]
  
      with tf.variable_scope("vae_internal", reuse=tf.AUTO_REUSE):
          # encoder
          enc_out = self.visual_encoder(inputs, clss, hparams, train)
          enc_out = tf.reshape(
              enc_out, [-1, 2 * hparams.bottleneck_bits]
          )
  
          # bottleneck
          sampled_bottleneck, b_loss = self.bottleneck(enc_out)
          losses = {"bottleneck_kl": tf.reduce_mean(b_loss)}
  
          if "bottleneck" in features:
              if (
                  common_layers.shape_list(features["bottleneck"])[0]
                  == 0
              ):
                  # return bottleneck for interpolation
                  # set losses['training'] = 0 so top() isn't called on it
                  # potential todo: use losses dict so we have kl_loss here for non stop
                  # gradient models
                  return sampled_bottleneck, {"training": 0.0}
              else:
                  # we want to use the given bottleneck
                  sampled_bottleneck = features["bottleneck"]
  
          # finalize bottleneck
          unbottleneck = sampled_bottleneck
  
          # decoder.
          dec_out = self.visual_decoder(unbottleneck, clss, hparams)
  
          # calculate training loss here lol
          rec_loss = -dec_out.log_prob(inputs)
          elbo = tf.reduce_mean(-(b_loss + rec_loss))
          losses["rec_loss"] = tf.reduce_mean(rec_loss)
          losses["training"] = -elbo
  
          if (
              not hasattr(self, "summarized_imgs")
              and self._hparams.mode != tf_estimator.ModeKeys.PREDICT
          ):
              self.summarized_imgs = True
              with tf.name_scope(None), tf.name_scope(
                  "train" if train else "test"
              ):
                  tf.summary.image("rendered_out", dec_out.mean())
                  tf.summary.image("rendered_og", inputs)
  
      return dec_out.mean(), losses
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from magenta.models.image_stylization import ops
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 5 for calling the ImageVAE obj's underlying member method code with example usage: obj.vae_internal(features, hparams, train) and returns: dec_out,  losses, sampled_bottleneck,  dict