class_docstr: A Conv2DTranspose decoder to generate Atari images from a latent space.
  
  Wraps an initial single linear layer with a stack of 4 Conv2DTranspose layers (with
  layer normalization) and a diag Gaussian, from which we then sample the final image.
  Sampling is done with a fixed stddev=1.0 and using the mean values coming from the
  last Conv2DTranspose layer.
docstr: Performs a forward pass through the Conv2D transpose decoder.
  
  Args:
      h: The deterministic hidden state of the sequence model.
      z: The sequence of stochastic discrete representations of the original
          observation input. Note: `z` is not used for the dynamics predictor
          model (which predicts z from h).
docstr_tok: ["Performs","a","forward","pass","through","the","Conv2D","transpose","decoder",".","Args",":","h",":","The","deterministic","hidden","state","of","the","sequence","model",".","z",":","The","sequence","of","stochastic","discrete","representations","of","the","original","observation","input",".","Note",":","`","z","`","is","not","used","for","the","dynamics","predictor","model","(","which","predicts","z","from","h",")","."]
code: def call(self, h, z):
      """Performs a forward pass through the Conv2D transpose decoder.
  
      Args:
          h: The deterministic hidden state of the sequence model.
          z: The sequence of stochastic discrete representations of the original
              observation input. Note: `z` is not used for the dynamics predictor
              model (which predicts z from h).
      """
      # Flatten last two dims of z.
      assert len(z.shape) == 3
      z_shape = tf.shape(z)
      z = tf.reshape(z, shape=(z_shape[0], -1))
      assert len(z.shape) == 2
      input_ = tf.concat([h, z], axis=-1)
      input_.set_shape(
          [
              None,
              (
                  get_num_z_categoricals(self.model_size)
                  * get_num_z_classes(self.model_size)
                  + get_gru_units(self.model_size)
              ),
          ]
      )
  
      # Feed through initial dense layer to get the right number of input nodes
      # for the first conv2dtranspose layer.
      out = self.dense_layer(input_)
      # Reshape to image format.
      out = tf.reshape(out, shape=(-1,) + self.input_dims)
  
      # Pass through stack of Conv2DTransport layers (and layer norms).
      for conv_transpose_2d, layer_norm in zip(
          self.conv_transpose_layers, self.layer_normalizations
      ):
          out = tf.nn.silu(layer_norm(inputs=conv_transpose_2d(out)))
      # Last output conv2d-transpose layer:
      out = self.output_conv2d_transpose(out)
      out += 0.5  # See Danijar's code
      out_shape = tf.shape(out)
  
      # Interpret output as means of a diag-Gaussian with std=1.0:
      # From [2]:
      # "Distributions: The image predictor outputs the mean of a diagonal Gaussian
      # likelihood with unit variance, ..."
  
      # Reshape `out` for the diagonal multi-variate Gaussian (each pixel is its own
      # independent (b/c diagonal co-variance matrix) variable).
      loc = tf.reshape(out, shape=(out_shape[0], -1))
  
      return loc
  
file_path: ray/rllib/algorithms/dreamerv3/tf/models/components/conv_transpose_atari.py
imports: from typing import Optional
  import numpy
  from ray.rllib.algorithms.dreamerv3.utils import get_cnn_multiplier, get_gru_units, get_num_z_categoricals, get_num_z_classes
  from ray.rllib.utils.framework import try_import_tf
label_desc: Use image node_id 2 for calling the ConvTransposeAtari obj's underlying member method code with example usage: obj.call(h, z) and returns: loc