code: def build_trainable_highway_flow(
      width,
      residual_fraction_initial_value=0.5,
      activation_fn=None,
      gate_first_n=None,
      seed=None,
      validate_args=False,
  ):
      """Builds a HighwayFlow parameterized by trainable variables.
  
      The variables are transformed to enforce the following parameter constraints:
  
      - `residual_fraction` is bounded between 0 and 1.
      - `upper_diagonal_weights_matrix` is a randomly initialized (lower) diagonal
         matrix with positive diagonal of size `width x width`.
      - `lower_diagonal_weights_matrix` is a randomly initialized lower diagonal
         matrix with ones on the diagonal of size `width x width`;
      - `bias` is a randomly initialized vector of size `width`.
  
      Args:
        width: Input dimension of the bijector.
        residual_fraction_initial_value: Initial value for gating parameter, must be
          between 0 and 1.
        activation_fn: Callable invertible activation function
          (e.g., `tf.nn.softplus`), or `None`.
        gate_first_n: Decides which part of the input should be gated (useful for
          example when using auxiliary variables).
        seed: Seed for random initialization of the weights.
        validate_args: Python `bool`. Whether to validate input with runtime
            assertions.
            Default value: `False`.
  
      Returns:
        trainable_highway_flow: The initialized bijector.
      """
  
      residual_fraction_initial_value = tf.convert_to_tensor(
          residual_fraction_initial_value,
          dtype_hint=tf.float32,
          name="residual_fraction_initial_value",
      )
      dtype = residual_fraction_initial_value.dtype
  
      bias_seed, upper_seed, lower_seed = samplers.split_seed(seed, n=3)
      lower_bijector = chain.Chain(
          [
              transform_diagonal.TransformDiagonal(
                  diag_bijector=shift.Shift(1.0)
              ),
              pad.Pad(paddings=[(1, 0), (0, 1)]),
              fill_triangular.FillTriangular(),
          ]
      )
      unconstrained_lower_initial_values = samplers.normal(
          shape=lower_bijector.inverse_event_shape([width, width]),
          mean=0.0,
          stddev=0.01,
          seed=lower_seed,
      )
      upper_bijector = fill_scale_tril.FillScaleTriL(
          diag_bijector=softplus.Softplus(), diag_shift=None
      )
      unconstrained_upper_initial_values = samplers.normal(
          shape=upper_bijector.inverse_event_shape([width, width]),
          mean=0.0,
          stddev=0.01,
          seed=upper_seed,
      )
  
      return HighwayFlow(
          residual_fraction=deferred_tensor.TransformedVariable(
              initial_value=residual_fraction_initial_value,
              bijector=sigmoid.Sigmoid(),
              dtype=dtype,
          ),
          activation_fn=activation_fn,
          bias=tf.Variable(
              samplers.normal(
                  (width,), mean=0.0, stddev=0.01, seed=bias_seed
              ),
              dtype=dtype,
          ),
          upper_diagonal_weights_matrix=deferred_tensor.TransformedVariable(
              initial_value=upper_bijector.forward(
                  unconstrained_upper_initial_values
              ),
              bijector=upper_bijector,
              dtype=dtype,
          ),
          lower_diagonal_weights_matrix=deferred_tensor.TransformedVariable(
              initial_value=lower_bijector.forward(
                  unconstrained_lower_initial_values
              ),
              bijector=lower_bijector,
              dtype=dtype,
          ),
          gate_first_n=gate_first_n,
          validate_args=validate_args,
      )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import fill_scale_tril
  from tensorflow_probability.python.bijectors import fill_triangular
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import transform_diagonal
  from tensorflow_probability.python.experimental.bijectors import scalar_function_with_inferred_inverse
  from tensorflow_probability.python.internal import cache_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 1 for calling a global function with example usage: build_trainable_highway_flow(width, residual_fraction_initial_value, activation_fn, gate_first_n, seed, validate_args) and returns: HighwayFlow