code: def test_aaa_glow_training(self, depths, split_plans, prior_type):
      with tf.Graph().as_default():
          _, x_mask, _ = self.get_data()
          x = tf.random_normal(
              (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
              mean=10.0,
              stddev=3.0,
              dtype=DTYPE,
          )
          bias = common_attention.attention_bias_ignore_padding(
              1.0 - x_mask
          )
          hparams = self.get_hparams()
          hparams.prior_type = prior_type
          hparams.depths = depths
          hparams.split_plans = split_plans
          n_levels = len(hparams.depths.split("/"))
          kwargs = self.get_kwargs(x_mask, hparams)
          _ = kwargs.pop("decoder_self_attention_bias")
  
          x_inv, _, _, _ = glow.glow(
              "glow",
              x,
              x_mask,
              bias,
              inverse=False,
              init=True,
              disable_dropout=True,
              **kwargs,
          )
          curr_dir = tempfile.mkdtemp()
          model_path = os.path.join(curr_dir, "model")
  
          with tf.Session() as session:
              saver = tf.train.Saver()
              session.run(tf.global_variables_initializer())
              session.run(x_inv)
              saver.save(session, model_path)
  
      with tf.Graph().as_default():
          _, x_mask, _ = self.get_data()
          x = tf.random_normal(
              (BATCH_SIZE, TARGET_LENGTH, N_CHANNELS),
              mean=10.0,
              stddev=3.0,
              dtype=DTYPE,
          )
          bias = common_attention.attention_bias_ignore_padding(
              1.0 - x_mask
          )
          hparams = self.get_hparams()
          hparams.depths = depths
          hparams.split_plans = split_plans
          kwargs = self.get_kwargs(x_mask, hparams)
          _ = kwargs.pop("decoder_self_attention_bias")
          log_q_z = gops.standard_normal_density(x, x_mask)
          log_q_z = tf.reduce_sum(log_q_z) / tf.reduce_sum(x_mask)
  
          x_inv, logabsdets, log_ps, zs = glow.glow(
              "glow",
              x,
              x_mask,
              bias,
              inverse=False,
              init=False,
              disable_dropout=True,
              **kwargs,
          )
          x_inv_inv, logabsdets_inv, log_ps_inv, _ = glow.glow(
              "glow",
              x_inv,
              x_mask,
              bias,
              inverse=True,
              split_zs=zs,
              init=False,
              disable_dropout=True,
              **kwargs,
          )
          logabsdets = tf.reduce_sum(
              logabsdets, axis=0
          ) / tf.reduce_sum(x_mask)
          logabsdets_inv = tf.reduce_sum(
              logabsdets_inv, axis=0
          ) / tf.reduce_sum(x_mask)
          log_ps = tf.reduce_sum(log_ps, axis=0) / tf.reduce_sum(x_mask)
          log_ps_inv = tf.reduce_sum(
              log_ps_inv, axis=0
          ) / tf.reduce_sum(x_mask)
  
          with tf.Session() as session:
              saver = tf.train.Saver()
              saver.restore(session, model_path)
              (
                  x,
                  x_inv,
                  x_inv_inv,
                  log_q_z,
                  logabsdets,
                  log_ps,
                  logabsdets_inv,
                  log_ps_inv,
              ) = session.run(
                  [
                      x,
                      x_inv,
                      x_inv_inv,
                      log_q_z,
                      logabsdets,
                      log_ps,
                      logabsdets_inv,
                      log_ps_inv,
                  ]
              )
              diff = x - x_inv_inv
              log_ps_diff = log_ps - log_ps_inv
              logabsdets_sum = logabsdets + logabsdets_inv
              self.assertEqual(
                  x_inv.shape,
                  (
                      BATCH_SIZE,
                      TARGET_LENGTH // (2 ** (n_levels - 1)),
                      N_CHANNELS,
                  ),
              )
              print(np.max(np.abs(diff)))
              print(np.max(np.abs(log_ps_diff)))
              print(np.max(np.abs(logabsdets_sum)))
              self.assertTrue(
                  np.allclose(diff, 0.0, atol=1e-4),
                  msg=np.max(np.abs(diff)),
              )
              self.assertTrue(
                  np.allclose(log_ps_diff, 0.0, atol=1e-4),
                  msg=np.max(np.abs(log_ps_diff)),
              )
              self.assertTrue(
                  np.allclose(logabsdets_sum, 0.0, atol=1e-4),
                  msg=np.max(np.abs(logabsdets_sum)),
              )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import tempfile
  from absl.testing import parameterized
  import numpy
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import transformer_glow_layers
  from tensor2tensor.layers import transformer_glow_layers_ops
  from tensor2tensor.models import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 10 for calling the TransformerGlowLayersTest obj's underlying member method code with example usage: obj.test_aaa_glow_training(depths, split_plans, prior_type) without return types