def call(self, inputs):
    """Runs the model to generate a distribution p(x_t | z_t, f).

    Args:
      inputs: A tuple of (z_{1:T}, f), where `z_{1:T}` is a tensor of
        shape [..., batch_size, timesteps, latent_size_dynamic], and `f`
        is of shape [..., batch_size, latent_size_static].

    Returns:
      A batched Independent distribution wrapping a set of Normal
      distributions over the pixels of x_t, where the Independent
      distribution has event shape [height, width, channels], batch
      shape [batch_size, timesteps], and sample shape [sample_shape,
      batch_size, timesteps, height, width, channels].
    """
    # We explicitly broadcast f to the same shape as z other than the final
    # dimension, because `tf.concat` can't automatically do this.
    dynamic, static = inputs
    timesteps = tf.shape(input=dynamic)[-2]
    static = static[..., tf.newaxis, :] + tf.zeros([timesteps, 1])
    latents = tf.concat([dynamic, static], axis=-1)  # (sample, N, T, latents)
    out = self.dense(latents)
    out = tf.reshape(out, (-1, 1, 1, self.hidden_size))
    out = self.conv_transpose1(out)
    out = self.conv_transpose2(out)
    out = self.conv_transpose3(out)
    out = self.conv_transpose4(out)  # (sample*N*T, h, w, c)
    expanded_shape = tf.concat(
        (tf.shape(input=latents)[:-1], tf.shape(input=out)[1:]), axis=0)
    out = tf.reshape(out, expanded_shape)  # (sample, N, T, h, w, c)
    return tfd.Independent(
        distribution=tfd.Normal(loc=out, scale=1.),
        reinterpreted_batch_ndims=3,  # wrap (h, w, c)
        name="decoded_image")
def call(self, inputs):
    """Runs the model to generate a distribution p(x_t | z_t, f).

    Args:
      inputs: A tuple of (z_{1:T}, f), where `z_{1:T}` is a tensor of
        shape [..., batch_size, timesteps, latent_size_dynamic], and `f`
        is of shape [..., batch_size, latent_size_static].

    Returns:
      A batched Independent distribution wrapping a set of Normal
      distributions over the pixels of x_t, where the Independent
      distribution has event shape [height, width, channels], batch
      shape [batch_size, timesteps], and sample shape [sample_shape,
      batch_size, timesteps, height, width, channels].
    """
    # We explicitly broadcast f to the same shape as z other than the final
    # dimension, because `tf.concat` can't automatically do this.
    dynamic, static = inputs
    timesteps = tf.shape(input=dynamic)[-2]
    static = static[..., tf.newaxis, :] + tf.zeros([timesteps, 1])
    latents = tf.concat([dynamic, static], axis=-1)  # (sample, N, T, latents)
    out = self.dense(latents)
    out = tf.reshape(out, (-1, 1, 1, self.hidden_size))
    out = self.conv_transpose1(out)
    out = self.conv_transpose2(out)
    out = self.conv_transpose3(out)
    out = self.conv_transpose4(out)  # (sample*N*T, h, w, c)
    expanded_shape = tf.concat(
        (tf.shape(input=latents)[:-1], tf.shape(input=out)[1:]), axis=0)
    out = tf.reshape(out, expanded_shape)  # (sample, N, T, h, w, c)
    return tfd.Independent(
        distribution=tfd.Normal(loc=out, scale=1.),
        reinterpreted_batch_ndims=3,  # wrap (h, w, c)
        name="decoded_image")
def call ( self , inputs ) : # We explicitly broadcast f to the same shape as z other than the final # dimension, because `tf.concat` can't automatically do this. dynamic , static = inputs timesteps = tf . shape ( input = dynamic ) [ - 2 ] static = static [ ... , tf . newaxis , : ] + tf . zeros ( [ timesteps , 1 ] ) latents = tf . concat ( [ dynamic , static ] , axis = - 1 ) # (sample, N, T, latents) out = self . dense ( latents ) out = tf . reshape ( out , ( - 1 , 1 , 1 , self . hidden_size ) ) out = self . conv_transpose1 ( out ) out = self . conv_transpose2 ( out ) out = self . conv_transpose3 ( out ) out = self . conv_transpose4 ( out ) # (sample*N*T, h, w, c) expanded_shape = tf . concat ( ( tf . shape ( input = latents ) [ : - 1 ] , tf . shape ( input = out ) [ 1 : ] ) , axis = 0 ) out = tf . reshape ( out , expanded_shape ) # (sample, N, T, h, w, c) return tfd . Independent ( distribution = tfd . Normal ( loc = out , scale = 1. ) , reinterpreted_batch_ndims = 3 , # wrap (h, w, c) name = "decoded_image" )
Runs the model to generate a distribution p(x_t | z_t, f).

    Args:
      inputs: A tuple of (z_{1:T}, f), where `z_{1:T}` is a tensor of
        shape [..., batch_size, timesteps, latent_size_dynamic], and `f`
        is of shape [..., batch_size, latent_size_static].

    Returns:
      A batched Independent distribution wrapping a set of Normal
      distributions over the pixels of x_t, where the Independent
      distribution has event shape [height, width, channels], batch
      shape [batch_size, timesteps], and sample shape [sample_shape,
      batch_size, timesteps, height, width, channels].
Runs the model to generate a distribution p ( x_t | z_t f ) .