code: def flow_step_glow(
      name, x, x_mask, split_dims, inverse, init, dtype, **kwargs
  ):
      """One step of flow."""
      conv_fn = multihead_invertible_1x1_conv_np
      with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
          reversible_ops = []
          for _, split_dim in enumerate(split_dims):
              identity_first = True
              reversible_ops += [
                  functools.partial(actnorm, name="actnorm", init=init)
              ]
              if split_dim in "ca":
                  multihead_split = "a" if split_dim == "c" else "c"
                  reversible_ops += [
                      functools.partial(
                          conv_fn,
                          name="conv_{}".format(multihead_split),
                          multihead_split=multihead_split,
                          dtype=dtype,
                      )
                  ]
              reversible_ops += [
                  functools.partial(
                      coupling,
                      name="coupling_{}".format(split_dim),
                      split_dim=split_dim,
                      identity_first=identity_first,
                      init=init,
                      **kwargs,
                  )
              ]
          if inverse:
              reversible_ops = reversible_ops[::-1]
  
          logabsdets = tf.constant(0.0, dtype=dtype)
          for reversible_op in reversible_ops:
              x, logabsdet = reversible_op(
                  x=x, x_mask=x_mask, inverse=inverse
              )
              logabsdets += logabsdet
          return x, logabsdets
  
file_path: tensor2tensor/tensor2tensor/layers/transformer_glow_layers.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import numpy
  import scipy
  from tensor2tensor.layers import common_layers
  import tensor2tensor.layers.transformer_glow_layers_ops
  import tensorflow.compat.v1
label_desc: Use image node_id 6 for calling a global function with example usage: flow_step_glow(name, x, x_mask, split_dims, inverse, init, dtype) and returns: x,  logabsdets