code: def multihead_invertible_1x1_conv_np(
      name, x, x_mask, multihead_split, inverse, dtype
  ):
      """Multi-head 1X1 convolution on x."""
      batch_size, length, n_channels_all = common_layers.shape_list(x)
      assert n_channels_all % 32 == 0
      n_channels = 32
      n_1x1_heads = n_channels_all // n_channels
  
      def get_init_np():
          """Initializer function for multihead 1x1 parameters using numpy."""
          results = []
          for _ in range(n_1x1_heads):
              random_matrix = np.random.rand(n_channels, n_channels)
              np_w = scipy.linalg.qr(random_matrix)[0].astype("float32")
              np_p, np_l, np_u = scipy.linalg.lu(np_w)
              np_s = np.diag(np_u)
              np_sign_s = np.sign(np_s)[np.newaxis, :]
              np_log_s = np.log(np.abs(np_s))[np.newaxis, :]
              np_u = np.triu(np_u, k=1)
              results.append(
                  np.concatenate(
                      [np_p, np_l, np_u, np_sign_s, np_log_s], axis=0
                  )
              )
          return tf.convert_to_tensor(np.stack(results, axis=0))
  
      def get_mask_init():
          ones = tf.ones(
              [n_1x1_heads, n_channels, n_channels], dtype=dtype
          )
          l_mask = tf.matrix_band_part(
              ones, -1, 0
          ) - tf.matrix_band_part(ones, 0, 0)
          u_mask = tf.matrix_band_part(
              ones, 0, -1
          ) - tf.matrix_band_part(ones, 0, 0)
          return tf.stack([l_mask, u_mask], axis=0)
  
      with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
          params = tf.get_variable(
              "params", initializer=get_init_np, dtype=dtype
          )
          mask_params = tf.get_variable(
              "mask_params",
              initializer=get_mask_init,
              dtype=dtype,
              trainable=False,
          )
  
          p = tf.stop_gradient(params[:, :n_channels, :])
          l = params[:, n_channels : 2 * n_channels, :]
          u = params[:, 2 * n_channels : 3 * n_channels, :]
          sign_s = tf.stop_gradient(params[:, 3 * n_channels, :])
          log_s = params[:, 3 * n_channels + 1, :]
  
          l_mask = mask_params[0]
          u_mask = mask_params[1]
  
          l_diag = l * l_mask + (
              tf.eye(n_channels, n_channels, [n_1x1_heads], dtype=dtype)
          )
          u_diag = u * u_mask + (tf.matrix_diag(sign_s * tf.exp(log_s)))
          w = tf.matmul(p, tf.matmul(l_diag, u_diag))
  
          if multihead_split == "a":
              x = tf.reshape(
                  x, [batch_size, length, n_channels, n_1x1_heads]
              )
              x = tf.transpose(x, [3, 0, 1, 2])
          elif multihead_split == "c":
              x = tf.reshape(
                  x, [batch_size, length, n_1x1_heads, n_channels]
              )
              x = tf.transpose(x, [2, 0, 1, 3])
          else:
              raise ValueError("Multihead split not supported.")
          # [n_1x1_heads, batch_size, length, n_channels]
  
          if not inverse:
              # [n_1x1_heads, 1, n_channels, n_channels]
              x = tf.matmul(x, w[:, tf.newaxis, :, :])
          else:
              w_inv = tf.matrix_inverse(w)
              x = tf.matmul(x, w_inv[:, tf.newaxis, :, :])
  
          if multihead_split == "a":
              x = tf.transpose(x, [1, 2, 3, 0])
              x = tf.reshape(
                  x, [batch_size, length, n_channels * n_1x1_heads]
              )
          elif multihead_split == "c":
              x = tf.transpose(x, [1, 2, 0, 3])
              x = tf.reshape(
                  x, [batch_size, length, n_1x1_heads * n_channels]
              )
          else:
              raise ValueError("Multihead split not supported.")
  
          x_length = tf.reduce_sum(x_mask, -1)
          logabsdet = x_length * tf.reduce_sum(log_s)
          if inverse:
              logabsdet *= -1
      return x, logabsdet
  
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 2 for calling a global function with example usage: multihead_invertible_1x1_conv_np(name, x, x_mask, multihead_split, inverse, dtype) and returns: x,  logabsdet, tf, tf