code: def _make_gated_bijector_fn():
      def _bijector_fn(x, output_units):
          if tensorshape_util.rank(x.shape) == 1:
              x = x[tf.newaxis, ...]
              reshape_output = lambda x: x[0]
          else:
              reshape_output = lambda x: x
  
          out = tf_keras.tf1_layers.dense(
              inputs=x, units=2 * output_units
          )
          shift, logit_gate = tf.split(out, 2, axis=-1)
          shift = reshape_output(shift)
          logit_gate = reshape_output(logit_gate)
          gate = tf.nn.sigmoid(logit_gate)
          return shift_lib.Shift(shift=(1.0 - gate) * shift)(
              scale.Scale(scale=gate)
          )
  
      return tf1.make_template("gated_bijector", _bijector_fn)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import inline
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import real_nvp
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 1 for calling a global function with example usage: _make_gated_bijector_fn() and returns: tf1, shift_lib