class_docstr: Disentangled Sequential Variational Autoencoder.
  
  The disentangled sequential variational autoencoder posits a generative
  model in which a static, time-invariant latent variable `f` is sampled
  from a prior `p(f)`, a dynamic, time-variant latent variable `z_t` at
  timestep `t` is sampled from a conditional distribution
  `p(z_t | z_{<t})`, and an observation `x_t` is generated by a
  probabilistic decoder `p(x_t | z_t, f)`. The full generative model is
  defined as
  
  ```none
  p(x_{1:T}, z_{1:T}, f) = p(f) prod_{t=1}^T p(z_t | z_{<t}) p(x_t | z_t, f).
  ```
  
  We then posit an approximate posterior over the latent variables in the
  form of a probabilistic encoder `q(z_{1:T}, f | x_{1:T})`. Paired with
  the probabilistic decoder, we can form a sequential variational
  autoencoder model. Variational inference can be used to fit the model by
  decomposing the log marginal distribution `log p(x_{1:T})` into the
  evidence lower bound (ELBO) and the KL divergence between the true and
  approximate posteriors over the latent variables
  
  ```none
  log p(x) = -KL[q(z_{1:T},f|x_{1:T}) || p(x_{1:T},z_{1:T},f)]
             + KL[q(z_{1:T},f|x_{1:T}) || p(z_{1:T},f|x_{1:T})]
           = ELBO + KL[q(z_{1:T},f|x_{1:T}) || p(z_{1:T},f|x_{1:T})]
          >= ELBO  # Jensen's inequality for KL divergence.
          >= int int q(z_{1:T},f|x_{1:T}) [
              log p(x_{1:T},z_{1:T},f) - log q(z_{1:T},f|x_{1:T}) ] dz_{1:T} df.
  ```
  
  We then maximize the ELBO with respect to the model's parameters.
  
  The approximate posterior `q(z_{1:T}, f | x_{1:T})` can be formulated in
  two ways. The first formulation is a distribution that factorizes
  across timesteps,
  
  ```none
  q(z_{1:T}, f | x_{1:T}) = q(f | x_{1:T}) prod_{t=1}^T q(z_t | x_t),
  ```
  
  where `q(f | x_{1:T})` is a multivariate Gaussian parameterized by a
  bidirectional LSTM-based model, and `q(z_t | x_t)` is a multivariate
  Gaussian parameterized by a convolutional model. This is known as the
  "factorized" `q` distribution.
  
  The second formulation is a distribution
  
  ```none
  q(z_{1:T}, f | x_{1:T}) = q(f | x_{1:T}) q(z_{1:T} | f, x_{1:T}),
  ```
  
  where `q(z_{1:T} | f, x_{1:T})` is a multivariate Gaussian parameterized
  by a model consisting of a bidirectional LSTM followed by a basic RNN,
  and `q(f | x_{1:T})` is the same as previously described. This is known
  as the "full" `q` distribution.
class_docstr_tok: ["Disentangled","Sequential","Variational","Autoencoder",".","The","disentangled","sequential","variational","autoencoder","posits","a","generative","model","in","which","a","static",",","time-invariant","latent","variable","`","f","`","is","sampled","from","a","prior","`","p","(","f",")","`",",","a","dynamic",",","time-variant","latent","variable","`","z_t","`","at","timestep","`","t","`","is","sampled","from","a","conditional","distribution","`","p","(","z_t","|","z_","{","<","t","}",")","`",",","and","an","observation","`","x_t","`","is","generated","by","a","probabilistic","decoder","`","p","(","x_t","|","z_t",",","f",")","`",".","The","full","generative","model","is","defined","as","``","`","none","p","(","x_","{","1",":","T","}",",","z_","{","1",":","T","}",",","f",")","=","p","(","f",")","prod_","{","t=1","}","^T","p","(","z_t","|","z_","{","<","t","}",")","p","(","x_t","|","z_t",",","f",")",".","``","`","We","then","posit","an","approximate","posterior","over","the","latent","variables","in","the","form","of","a","probabilistic","encoder","`","q","(","z_","{","1",":","T","}",",","f","|","x_","{","1",":","T","}",")","`",".","Paired","with","the","probabilistic","decoder",",","we","can","form","a","sequential","variational","autoencoder","model",".","Variational","inference","can","be","used","to","fit","the","model","by","decomposing","the","log","marginal","distribution","`","log","p","(","x_","{","1",":","T","}",")","`","into","the","evidence","lower","bound","(","ELBO",")","and","the","KL","divergence","between","the","true","and","approximate","posteriors","over","the","latent","variables","``","`","none","log","p","(","x",")","=","-KL","[","q","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","||","p","(","x_","{","1",":","T","}",",","z_","{","1",":","T","}",",","f",")","]","+","KL","[","q","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","||","p","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","]","=","ELBO","+","KL","[","q","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","||","p","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","]",">","=","ELBO","#","Jensen","'s","inequality","for","KL","divergence",".",">","=","int","int","q","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","[","log","p","(","x_","{","1",":","T","}",",","z_","{","1",":","T","}",",","f",")","-","log","q","(","z_","{","1",":","T","}",",","f|x_","{","1",":","T","}",")","]","dz_","{","1",":","T","}","df",".","``","`","We","then","maximize","the","ELBO","with","respect","to","the","model","'s","parameters",".","The","approximate","posterior","`","q","(","z_","{","1",":","T","}",",","f","|","x_","{","1",":","T","}",")","`","can","be","formulated","in","two","ways",".","The","first","formulation","is","a","distribution","that","factorizes","across","timesteps",",","``","`","none","q","(","z_","{","1",":","T","}",",","f","|","x_","{","1",":","T","}",")","=","q","(","f","|","x_","{","1",":","T","}",")","prod_","{","t=1","}","^T","q","(","z_t","|","x_t",")",",","``","`","where","`","q","(","f","|","x_","{","1",":","T","}",")","`","is","a","multivariate","Gaussian","parameterized","by","a","bidirectional","LSTM-based","model",",","and","`","q","(","z_t","|","x_t",")","`","is","a","multivariate","Gaussian","parameterized","by","a","convolutional","model",".","This","is","known","as","the","''","factorized","''","`","q","`","distribution",".","The","second","formulation","is","a","distribution","``","`","none","q","(","z_","{","1",":","T","}",",","f","|","x_","{","1",":","T","}",")","=","q","(","f","|","x_","{","1",":","T","}",")","q","(","z_","{","1",":","T","}","|","f",",","x_","{","1",":","T","}",")",",","``","`","where","`","q","(","z_","{","1",":","T","}","|","f",",","x_","{","1",":","T","}",")","`","is","a","multivariate","Gaussian","parameterized","by","a","model","consisting","of","a","bidirectional","LSTM","followed","by","a","basic","RNN",",","and","`","q","(","f","|","x_","{","1",":","T","}",")","`","is","the","same","as","previously","described",".","This","is","known","as","the","``","full","''","`","q","`","distribution","."]
docstr: Constructs a Disentangled Sequential Variational Autoencoder.
  
  Args:
    latent_size_static: Integer dimensionality of the static,
      time-invariant latent variable `f`.
    latent_size_dynamic: Integer dimensionality of each dynamic,
      time-variant latent variable `z_t`.
    hidden_size: Integer dimensionality of the model intermediates.
    channels: Integer depth of the output of the decoder.
    latent_posterior: Either "factorized" or "full" to indicate the
      formulation for the latent posterior `q`.
code: def __init__(
      self,
      latent_size_static,
      latent_size_dynamic,
      hidden_size,
      channels,
      latent_posterior,
  ):
      """Constructs a Disentangled Sequential Variational Autoencoder.
  
      Args:
        latent_size_static: Integer dimensionality of the static,
          time-invariant latent variable `f`.
        latent_size_dynamic: Integer dimensionality of each dynamic,
          time-variant latent variable `z_t`.
        hidden_size: Integer dimensionality of the model intermediates.
        channels: Integer depth of the output of the decoder.
        latent_posterior: Either "factorized" or "full" to indicate the
          formulation for the latent posterior `q`.
      """
      super(DisentangledSequentialVAE, self).__init__()
      self.latent_size_static = latent_size_static
      self.latent_size_dynamic = latent_size_dynamic
      self.hidden_size = hidden_size
      self.channels = channels
      self.latent_posterior = latent_posterior
      self.static_prior = LearnableMultivariateNormalDiag(
          latent_size_static
      )
      self.dynamic_prior = LearnableMultivariateNormalDiagCell(
          latent_size_dynamic, hidden_size
      )
      self.decoder = Decoder(hidden_size, channels)
      self.compressor = Compressor(hidden_size)
      self.static_encoder = EncoderStatic(
          latent_size_static, hidden_size
      )
      if latent_posterior == "factorized":
          self.dynamic_encoder = EncoderDynamicFactorized(
              latent_size_dynamic, hidden_size
          )
      else:
          self.dynamic_encoder = EncoderDynamicFull(
              latent_size_dynamic, hidden_size
          )
  
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 1 to create a new DisentangledSequentialVAE object from inherited base classes: tf_keras.v1 with example: obj = DisentangledSequentialVAE(latent_size_static, latent_size_dynamic, hidden_size, channels, latent_posterior)