class_docstr: Probabilistic encoder for the time-variant latent variable `z_t`.
  
  The conditional distribution `q(z_{1:T} | x_{1:T}, f)` is a
  multivariate normal distribution on `R^{latent_size}` at each timestep
  `t`, conditioned on both an intermediate representation of the inputs
  `x_t` from the convolutional encoder, and on a sample of the static
  latent variable `f` at each timestep. The parameters are computed by
  passing the inputs through a bidirectional LSTM function, then passing
  these intermediates through an RNN function and an affine function to
  yield normal parameters for `q(z_t | x_{1:T}, f)`.
  
  In this formulation, we posit that `z_t` is conditionally dependent on
  `f`.
  
  Together with the EncoderStatic class, we can formulate the full
  approximate later posterior `q` inference ("encoder") model as
  
  ```none
  q(z_{1:T}, f | x_{1:T}) = q(f | x_{1:T}) q(z_{1:T} | f, x_{1:T}).
  ```
class_docstr_tok: ["Probabilistic","encoder","for","the","time-variant","latent","variable","`","z_t","`",".","The","conditional","distribution","`","q","(","z_","{","1",":","T","}","|","x_","{","1",":","T","}",",","f",")","`","is","a","multivariate","normal","distribution","on","`","R^","{","latent_size","}","`","at","each","timestep","`","t","`",",","conditioned","on","both","an","intermediate","representation","of","the","inputs","`","x_t","`","from","the","convolutional","encoder",",","and","on","a","sample","of","the","static","latent","variable","`","f","`","at","each","timestep",".","The","parameters","are","computed","by","passing","the","inputs","through","a","bidirectional","LSTM","function",",","then","passing","these","intermediates","through","an","RNN","function","and","an","affine","function","to","yield","normal","parameters","for","`","q","(","z_t","|","x_","{","1",":","T","}",",","f",")","`",".","In","this","formulation",",","we","posit","that","`","z_t","`","is","conditionally","dependent","on","`","f","`",".","Together","with","the","EncoderStatic","class",",","we","can","formulate","the","full","approximate","later","posterior","`","q","`","inference","(","``","encoder","''",")","model","as","``","`","none","q","(","z_","{","1",":","T","}",",","f","|","x_","{","1",":","T","}",")","=","q","(","f","|","x_","{","1",":","T","}",")","q","(","z_","{","1",":","T","}","|","f",",","x_","{","1",":","T","}",")",".","``","`"]
docstr: Runs the model to generate a distribution `q(z_{1:T} | x_{1:T}, f)`.
  
  This generates a list of batched MultivariateNormalDiag
  distributions using the output of the recurrent model at each
  timestep to parameterize each distribution.
  
  Args:
    inputs: A tuple of a batch of intermediate representations of
      image frames across all timesteps of shape [..., batch_size,
      timesteps, dimensions], and a sample of the static latent
      variable `f` of shape [..., batch_size, latent_size].
  
  Returns:
    A batch of MultivariateNormalDiag distributions with event shape
    [latent_size], batch shape [broadcasted_shape, batch_size,
    timesteps], and sample shape [sample_shape, broadcasted_shape,
    batch_size, timesteps, latent_size], where `broadcasted_shape` is
    the broadcasted sampled shape between the inputs and static
    sample.
code: def call(self, inputs):
      """Runs the model to generate a distribution `q(z_{1:T} | x_{1:T}, f)`.
  
      This generates a list of batched MultivariateNormalDiag
      distributions using the output of the recurrent model at each
      timestep to parameterize each distribution.
  
      Args:
        inputs: A tuple of a batch of intermediate representations of
          image frames across all timesteps of shape [..., batch_size,
          timesteps, dimensions], and a sample of the static latent
          variable `f` of shape [..., batch_size, latent_size].
  
      Returns:
        A batch of MultivariateNormalDiag distributions with event shape
        [latent_size], batch shape [broadcasted_shape, batch_size,
        timesteps], and sample shape [sample_shape, broadcasted_shape,
        batch_size, timesteps, latent_size], where `broadcasted_shape` is
        the broadcasted sampled shape between the inputs and static
        sample.
      """
      # We explicitly broadcast `x` and `f` to the same shape other than the final
      # dimension, because `tf1.concat` can't automatically do this. This will
      # entail adding a `timesteps` dimension to `f` to give the shape `(...,
      # batch, timesteps, latent)`, and then broadcasting the sample shapes of
      # both tensors to the same shape.
      features, static_sample = inputs
      length = tf1.shape(input=features)[-2]
      static_sample = static_sample[..., tf1.newaxis, :] + tf1.zeros(
          [length, 1]
      )
      sample_shape_static = tf1.shape(input=static_sample)[:-3]
      sample_shape_inputs = tf1.shape(input=features)[:-3]
      broadcast_shape_inputs = tf1.concat(
          (sample_shape_static, [1, 1, 1]), 0
      )
      broadcast_shape_static = tf1.concat(
          (sample_shape_inputs, [1, 1, 1]), 0
      )
      features = features + tf1.zeros(broadcast_shape_inputs)
      static_sample = static_sample + tf1.zeros(broadcast_shape_static)
      # `combined` will have shape (..., batch, T, hidden+latent).
      combined = tf1.concat((features, static_sample), axis=-1)
      # TODO(dusenberrymw): Remove these reshaping commands after b/113126249 is
      # fixed.
      collapsed_shape = tf1.concat(
          ([-1], tf1.shape(input=combined)[-2:]), axis=0
      )
      out = tf1.reshape(combined, collapsed_shape)
      out = self.bilstm(out)  # (sample*batch, T, hidden_size)
      out = self.rnn(out)  # (sample*batch, T, hidden_size)
      expanded_shape = tf1.concat(
          (tf1.shape(input=combined)[:-2], tf1.shape(input=out)[1:]),
          axis=0,
      )
      out = tf1.reshape(
          out, expanded_shape
      )  # (sample, batch, T, hidden_size)
      out = self.output_layer(out)  # (sample, batch, T, 2*latent_size)
      loc = out[..., : self.latent_size]
      scale_diag = (
          tf1.nn.softplus(out[..., self.latent_size :]) + 1e-5
      )  # keep > 0
      return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
  
imports: from datetime import datetime
  import functools
  from absl import app
  from absl import flags
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.examples import sprites_dataset
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 2 for calling the EncoderDynamicFull obj's underlying member method code with example usage: obj.call(inputs) and returns: tfd