code: def main(argv):
      del argv  # unused
  
      tf1.enable_eager_execution()
      tf1.set_random_seed(FLAGS.seed)
      timestamp = datetime.strftime(datetime.today(), "%y%m%d_%H%M%S")
      FLAGS.logdir = FLAGS.logdir.format(timestamp=timestamp)
      FLAGS.model_dir = FLAGS.model_dir.format(timestamp=timestamp)
      if not tf1.io.gfile.exists(FLAGS.model_dir):
          tf1.io.gfile.makedirs(FLAGS.model_dir)
  
      sprites_data = sprites_dataset.SpritesDataset(
          fake_data=FLAGS.fake_data
      )
  
      model = DisentangledSequentialVAE(
          latent_size_static=FLAGS.latent_size_static,
          latent_size_dynamic=FLAGS.latent_size_dynamic,
          hidden_size=FLAGS.hidden_size,
          channels=sprites_data.channels,
          latent_posterior=FLAGS.latent_posterior,
      )
  
      global_step = tf1.train.get_or_create_global_step()
      optimizer = tf1.train.AdamOptimizer(
          tf1.train.cosine_decay(
              FLAGS.learning_rate, global_step, FLAGS.max_steps
          )
      )
  
      checkpoint = tf1.train.Checkpoint(
          model=model, global_step=global_step, optimizer=optimizer
      )
      checkpoint_manager = tf1.train.CheckpointManager(
          checkpoint, directory=FLAGS.model_dir, max_to_keep=5
      )
      checkpoint.restore(checkpoint_manager.latest_checkpoint)
  
      writer = tf.summary.create_file_writer(FLAGS.logdir)
      writer.set_as_default()
  
      dataset = (
          sprites_data.train.map(lambda *x: x[0]).shuffle(1000).repeat()
      )
      dataset = dataset.batch(FLAGS.batch_size).take(FLAGS.max_steps)
  
      if FLAGS.enable_debug_logging:
          for inputs in dataset.prefetch(buffer_size=None):
              with tf.summary.record_if(
                  lambda: tf1.math.equal(
                      0, global_step % FLAGS.log_steps
                  )
              ):
                  tf.summary.histogram(
                      "image",
                      data=inputs,
                      step=tf1.train.get_or_create_global_step(),
                  )
  
              with tf1.GradientTape() as tape:
                  features = model.compressor(
                      inputs
                  )  # (batch, timesteps, hidden)
                  (
                      static_sample,
                      static_posterior,
                  ) = model.sample_static_posterior(
                      features, FLAGS.num_samples
                  )  # (samples, batch, latent)
                  (
                      dynamic_sample,
                      dynamic_posterior,
                  ) = model.sample_dynamic_posterior(
                      features, FLAGS.num_samples, static_sample
                  )  # (sampl, N, T, latent)
                  likelihood = model.decoder(
                      (dynamic_sample, static_sample)
                  )
  
                  reconstruction = tf1.reduce_mean(  # integrate samples
                      input_tensor=likelihood.mean()[
                          : FLAGS.num_reconstruction_samples
                      ],
                      axis=0,
                  )
                  visualize_reconstruction(
                      inputs,
                      reconstruction,
                      name="train_reconstruction",
                  )
  
                  static_prior = model.static_prior()
                  _, dynamic_prior = model.sample_dynamic_prior(
                      FLAGS.num_samples,
                      FLAGS.batch_size,
                      sprites_data.length,
                  )
  
                  if FLAGS.enable_debug_logging:
                      summarize_dist_params(
                          static_prior, "static_prior"
                      )
                      summarize_dist_params(
                          static_posterior, "static_posterior"
                      )
                      summarize_dist_params(
                          dynamic_prior, "dynamic_prior"
                      )
                      summarize_dist_params(
                          dynamic_posterior, "dynamic_posterior"
                      )
                      summarize_dist_params(likelihood, "likelihood")
  
                  static_prior_log_prob = static_prior.log_prob(
                      static_sample
                  )
                  static_posterior_log_prob = static_posterior.log_prob(
                      static_sample
                  )
                  dynamic_prior_log_prob = tf1.reduce_sum(
                      input_tensor=dynamic_prior.log_prob(
                          dynamic_sample
                      ),
                      axis=-1,
                  )  # sum time
                  dynamic_posterior_log_prob = tf1.reduce_sum(
                      input_tensor=dynamic_posterior.log_prob(
                          dynamic_sample
                      ),
                      axis=-1,
                  )  # sum time
                  likelihood_log_prob = tf1.reduce_sum(
                      input_tensor=likelihood.log_prob(inputs), axis=-1
                  )  # sum time
  
                  if FLAGS.enable_debug_logging:
                      with tf1.name_scope("log_probs"):
                          summarize_mean_in_nats_and_bits(
                              static_prior_log_prob,
                              FLAGS.latent_size_static,
                              "static_prior",
                          )
                          summarize_mean_in_nats_and_bits(
                              static_posterior_log_prob,
                              FLAGS.latent_size_static,
                              "static_posterior",
                          )
                          summarize_mean_in_nats_and_bits(
                              dynamic_prior_log_prob,
                              FLAGS.latent_size_dynamic
                              * sprites_data.length,
                              "dynamic_prior",
                          )
                          summarize_mean_in_nats_and_bits(
                              dynamic_posterior_log_prob,
                              FLAGS.latent_size_dynamic
                              * sprites_data.length,
                              "dynamic_posterior",
                          )
                          summarize_mean_in_nats_and_bits(
                              likelihood_log_prob,
                              sprites_data.frame_size**2
                              * sprites_data.channels
                              * sprites_data.length,
                              "likelihood",
                          )
  
                  elbo = tf1.reduce_mean(
                      input_tensor=static_prior_log_prob
                      - static_posterior_log_prob
                      + dynamic_prior_log_prob
                      - dynamic_posterior_log_prob
                      + likelihood_log_prob
                  )
                  loss = -elbo
                  tf.summary.scalar(
                      "elbo",
                      elbo,
                      step=tf1.train.get_or_create_global_step(),
                  )
  
              grads = tape.gradient(loss, model.variables)
              grads, global_norm = tf1.clip_by_global_norm(
                  grads, FLAGS.clip_norm
              )
              grads_and_vars = list(
                  zip(grads, model.variables)
              )  # allow reuse in py3
              if FLAGS.enable_debug_logging:
                  with tf1.name_scope("grads"):
                      tf.summary.scalar(
                          "global_norm_grads",
                          global_norm,
                          step=tf1.train.get_or_create_global_step(),
                      )
                      tf.summary.scalar(
                          "global_norm_grads_clipped",
                          tf1.linalg.global_norm(grads),
                          step=tf1.train.get_or_create_global_step(),
                      )
                  for grad, var in grads_and_vars:
                      with tf1.name_scope("grads"):
                          tf.summary.histogram(
                              "{}/grad".format(var.name),
                              data=grad,
                              step=tf1.train.get_or_create_global_step(),
                          )
                      with tf1.name_scope("vars"):
                          tf.summary.histogram(
                              var.name,
                              data=var,
                              step=tf1.train.get_or_create_global_step(),
                          )
              optimizer.apply_gradients(grads_and_vars, global_step)
  
          is_log_step = global_step.numpy() % FLAGS.log_steps == 0
          is_final_step = global_step.numpy() == FLAGS.max_steps
          if is_log_step or is_final_step:
              checkpoint_manager.save()
              print(
                  "ELBO ({}/{}): {}".format(
                      global_step.numpy(), FLAGS.max_steps, elbo.numpy()
                  )
              )
              with tf.summary.record_if(True):
                  val_data = sprites_data.test.take(20)
                  inputs = next(iter(val_data.shuffle(20).batch(3)))[0]
                  visualize_qualitative_analysis(
                      inputs, model, FLAGS.num_reconstruction_samples
                  )
  
      writer.flush()
  
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 6 for calling a global function with example usage: main(argv) without return types