code: def build_graph(
      device, input_shape, variable, num_inputs, axis, grad
  ):
      """Build a graph containing a sequence of concat operations.
  
      Args:
        device: string, the device to run on.
        input_shape: shape of the input tensors.
        variable: whether or not to randomize the input shape
        num_inputs: the number of inputs to concat
        axis: axis to be concat'ed
        grad: if True compute the gradient
  
      Returns:
        An array of tensors to run()
      """
      with ops.device("/%s:0" % device):
          if not variable:
              inputs = [
                  array_ops.zeros(input_shape)
                  for _ in range(num_inputs)
              ]
          else:
              if axis == 1:
                  inputs = [
                      array_ops.zeros(
                          [
                              input_shape[0],
                              random.randint(
                                  max(1, input_shape[1] - 5),
                                  input_shape[1] + 5,
                              ),
                          ]
                      )
                      for _ in range(num_inputs)
                  ]
              else:
                  inputs = [
                      array_ops.zeros(
                          [
                              random.randint(
                                  max(1, input_shape[0] - 5),
                                  input_shape[0] + 5,
                              ),
                              input_shape[1],
                          ]
                      )
                      for _ in range(num_inputs)
                  ]
  
          outputs = [array_ops.concat(inputs, axis) for _ in range(100)]
          if grad:
              return control_flow_ops.group(
                  *list(
                      itertools.chain.from_iterable(
                          [
                              gradients_impl.gradients(output, inputs)
                              for output in outputs
                          ]
                      )
                  )
              )
          else:
              return control_flow_ops.group(*outputs)
  
file_path: tensorflow-upstream/tensorflow/python/ops/concat_benchmark.py
imports: import itertools
  import random
  import time
  from tensorflow.core.protobuf import config_pb2
  from tensorflow.python.client import session
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import gradients_impl
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling a global function with example usage: build_graph(device, input_shape, variable, num_inputs, axis, grad) and returns: control_flow_ops, control_flow_ops