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: Reconstruct the given input sequences.
  
  Args:
    inputs: A batch of image sequences `x_{1:T}` of shape
      `[batch_size, timesteps, height, width, channels]`.
    samples: Number of samples to draw from the latent distributions.
    sample_static: Boolean for whether or not to randomly sample the
      static latent variable `f` from its prior distribution.
    sample_dynamic: Boolean for whether or not to randomly sample the
      dynamic latent variable `z_{1:T}` from its prior distribution.
    swap_static: Boolean for whether or not to swap the encodings for
      the static latent variable `f` between the examples.
    swap_dynamic: Boolean for whether or not to swap the encodings for
      the dynamic latent variable `z_{1:T}` between the examples.
    fix_static: Boolean for whether or not to share the same random
      sample of the static latent variable `f` from its prior across
      all examples.
    fix_dynamic: Boolean for whether or not to share the same random
      sample of the dynamic latent variable `z_{1:T}` from its prior
      across all examples.
  
  Returns:
    A batched Independent distribution wrapping a set of Normal
    distributions over the pixels of the reconstruction of the input,
    where the Independent distribution has event shape [height, width,
    channels], batch shape [samples, batch_size, timesteps], and
    sample shape [sample_shape, samples, batch_size, timesteps,
    height, width, channels].
code: def reconstruct(
      self,
      inputs,
      samples=1,
      sample_static=False,
      sample_dynamic=False,
      swap_static=False,
      swap_dynamic=False,
      fix_static=False,
      fix_dynamic=False,
  ):
      """Reconstruct the given input sequences.
  
      Args:
        inputs: A batch of image sequences `x_{1:T}` of shape
          `[batch_size, timesteps, height, width, channels]`.
        samples: Number of samples to draw from the latent distributions.
        sample_static: Boolean for whether or not to randomly sample the
          static latent variable `f` from its prior distribution.
        sample_dynamic: Boolean for whether or not to randomly sample the
          dynamic latent variable `z_{1:T}` from its prior distribution.
        swap_static: Boolean for whether or not to swap the encodings for
          the static latent variable `f` between the examples.
        swap_dynamic: Boolean for whether or not to swap the encodings for
          the dynamic latent variable `z_{1:T}` between the examples.
        fix_static: Boolean for whether or not to share the same random
          sample of the static latent variable `f` from its prior across
          all examples.
        fix_dynamic: Boolean for whether or not to share the same random
          sample of the dynamic latent variable `z_{1:T}` from its prior
          across all examples.
  
      Returns:
        A batched Independent distribution wrapping a set of Normal
        distributions over the pixels of the reconstruction of the input,
        where the Independent distribution has event shape [height, width,
        channels], batch shape [samples, batch_size, timesteps], and
        sample shape [sample_shape, samples, batch_size, timesteps,
        height, width, channels].
      """
      batch_size = tf1.shape(input=inputs)[-5]
      length = len(tf1.unstack(inputs, axis=-4))  # hack for graph mode
  
      features = self.compressor(
          inputs
      )  # (..., batch, timesteps, hidden)
  
      if sample_static:
          static_sample, _ = self.sample_static_prior(
              samples, batch_size, fix_static
          )
      else:
          static_sample, _ = self.sample_static_posterior(
              features, samples
          )
  
      if swap_static:
          static_sample = tf1.reverse(static_sample, axis=[1])
  
      if sample_dynamic:
          dynamic_sample, _ = self.sample_dynamic_prior(
              samples, batch_size, length, fix_dynamic
          )
      else:
          dynamic_sample, _ = self.sample_dynamic_posterior(
              features, samples, static_sample
          )
  
      if swap_dynamic:
          dynamic_sample = tf1.reverse(dynamic_sample, axis=[1])
  
      likelihood = self.decoder((dynamic_sample, static_sample))
      return likelihood
  
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 3 for calling the DisentangledSequentialVAE obj's underlying member method code with example usage: obj.reconstruct(inputs, samples, sample_static, sample_dynamic, swap_static, swap_dynamic, fix_static, fix_dynamic) and returns: likelihood