def execute_tpu_tf2(self, compute_fn, inputs):
    """Executes compute_fn on TPU with Tensorflow 2.X.

    Args:
      compute_fn: a function containing Tensorflow computation that takes a list
        of input numpy tensors, performs computation and returns output numpy
        tensors.
      inputs: a list of numpy arrays to feed input to the `compute_fn`.

    Returns:
      A list of numpy arrays or a single numpy array.
    """
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
    tf.config.experimental_connect_to_cluster(resolver)
    topology = tf.tpu.experimental.initialize_tpu_system(resolver)
    device_assignment = tf.tpu.experimental.DeviceAssignment.build(
        topology, num_replicas=1)
    strategy = tf.distribute.experimental.TPUStrategy(
        resolver, device_assignment=device_assignment)

    @tf.function
    def run():
      tf_inputs = [tf.constant(input_t) for input_t in inputs]
      return strategy.run(compute_fn, args=tf_inputs)
    outputs = run()
    tf.tpu.experimental.shutdown_tpu_system()
    return self.maybe_extract_single_output(outputs)
Executes compute_fn on TPU with Tensorflow 2.X.

    Args:
      compute_fn: a function containing Tensorflow computation that takes a list
        of input numpy tensors, performs computation and returns output numpy
        tensors.
      inputs: a list of numpy arrays to feed input to the `compute_fn`.

    Returns:
      A list of numpy arrays or a single numpy array.
    
inference_api/src/main/object_detection/utils/test_case.py