docs: def freeze_graph_tpu(model_path):
      """Custom freeze_graph implementation for Cloud TPU."""
  
      assert model_path
      assert FLAGS.tpu_name
      if FLAGS.tpu_name.startswith('grpc://'):
          tpu_grpc_url = FLAGS.tpu_name
      else:
          tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
              FLAGS.tpu_name, zone=None, project=None)
          tpu_grpc_url = tpu_cluster_resolver.get_master()
      sess = tf.Session(tpu_grpc_url)
  
      output_names = []
      with sess.graph.as_default():
          # Replicate the inference function for each TPU core.
          replicated_features = []
          for i in range(FLAGS.num_tpu_cores):
              features = tf.placeholder(
                  tf.float32, [None, go.N, go.N,
                               features_lib.NEW_FEATURES_PLANES],
                  name='pos_tensor_%d' % i)
              replicated_features.append((features,))
          outputs = tf.contrib.tpu.replicate(
              tpu_model_inference_fn, replicated_features)
  
          # The replicate op assigns names like output_0_shard_0 to the output
          # names. Give them human readable names.
          for i, (policy_output, value_output, _) in enumerate(outputs):
              policy_name = 'policy_output_%d' % i
              value_name = 'value_output_%d' % i
              output_names.extend([policy_name, value_name])
              tf.identity(policy_output, policy_name)
              tf.identity(value_output, value_name)
  
          tf.train.Saver().restore(sess, model_path)
  
      # Freeze the graph.
      model_def = tf.graph_util.convert_variables_to_constants(
          sess, sess.graph.as_graph_def(), output_names)
      with tf.gfile.GFile(model_path + '.pb', 'wb') as f:
          f.write(model_def.SerializeToString())