code: def test_tpu_estimator_summaries(
      self,
      use_tpu,
      want_loss,
      want_adanet_loss,
      want_eval_summary_loss,
      want_predictions,
  ):
      max_steps = 10
      config = tf.compat.v1.estimator.tpu.RunConfig(
          tf_random_seed=42,
          save_summary_steps=2,
          log_step_count_steps=max_steps,
      )
      assert config.log_step_count_steps
  
      def metric_fn(predictions):
          return {
              "predictions": tf_compat.v1.metrics.mean(
                  predictions["predictions"]
              )
          }
  
      estimator = TPUEstimator(
          head=make_regression_head(use_tpu),
          subnetwork_generator=SimpleGenerator(
              [_DNNBuilder("dnn", use_tpu=use_tpu)]
          ),
          max_iteration_steps=max_steps,
          model_dir=self.test_subdirectory,
          metric_fn=metric_fn,
          config=config,
          use_tpu=use_tpu,
          train_batch_size=64 if use_tpu else 0,
      )
      xor_features = [[1.0, 0.0], [0.0, 0], [0.0, 1.0], [1.0, 1.0]]
      xor_labels = [[1.0], [0.0], [1.0], [0.0]]
      train_input_fn = tu.dummy_input_fn(xor_features, xor_labels)
  
      estimator.train(input_fn=train_input_fn, max_steps=max_steps)
      eval_results = estimator.evaluate(
          input_fn=train_input_fn, steps=1
      )
      self.assertAlmostEqual(want_loss, eval_results["loss"], places=2)
      self.assertEqual(max_steps, eval_results["global_step"])
      self.assertEqual(0, eval_results["iteration"])
  
      subnetwork_subdir = os.path.join(
          self.test_subdirectory, "subnetwork/t0_dnn"
      )
  
      ensemble_subdir = os.path.join(
          self.test_subdirectory,
          "ensemble/t0_dnn_grow_complexity_regularized",
      )
  
      # TODO: Why is the adanet_loss written to 'loss'?
      self.assertAlmostEqual(
          want_adanet_loss,
          tu.check_eventfile_for_keyword(
              "loss", self.test_subdirectory
          ),
          places=1,
      )
      self.assertEqual(
          0.0,
          tu.check_eventfile_for_keyword(
              "iteration/adanet/iteration", self.test_subdirectory
          ),
      )
      self.assertAlmostEqual(
          3.0,
          tu.check_eventfile_for_keyword("scalar", subnetwork_subdir),
          places=3,
      )
      self.assertEqual(
          (3, 3, 1),
          tu.check_eventfile_for_keyword(
              # When TF 2 behavior is enabled AdaNet uses V2 summaries.
              "image"
              if tf_compat.is_v2_behavior_enabled()
              else "image/image/0",
              subnetwork_subdir,
          ),
      )
      self.assertAlmostEqual(
          5.0,
          tu.check_eventfile_for_keyword(
              "nested/scalar", subnetwork_subdir
          ),
          places=3,
      )
      self.assertAlmostEqual(
          want_adanet_loss,
          tu.check_eventfile_for_keyword(
              "adanet_loss/adanet/adanet_weighted_ensemble",
              ensemble_subdir,
          ),
          places=1,
      )
      self.assertAlmostEqual(
          0.0,
          tu.check_eventfile_for_keyword(
              "complexity_regularization/adanet/adanet_weighted_ensemble",
              ensemble_subdir,
          ),
          places=1,
      )
      self.assertAlmostEqual(
          1.0,
          tu.check_eventfile_for_keyword(
              "mixture_weight_norms/adanet/"
              "adanet_weighted_ensemble/subnetwork_0",
              ensemble_subdir,
          ),
          places=1,
      )
  
      # Eval metric summaries are always written out during eval.
      subnetwork_eval_subdir = os.path.join(subnetwork_subdir, "eval")
      self.assertAlmostEqual(
          want_eval_summary_loss,
          tu.check_eventfile_for_keyword(
              "loss", subnetwork_eval_subdir
          ),
          places=1,
      )
      # TODO: Check why some eval metrics are zero on TPU.
      self.assertAlmostEqual(
          0.0 if use_tpu else want_eval_summary_loss,
          tu.check_eventfile_for_keyword(
              "average_loss", subnetwork_eval_subdir
          ),
          places=1,
      )
      self.assertAlmostEqual(
          want_predictions,
          tu.check_eventfile_for_keyword(
              "predictions", subnetwork_eval_subdir
          ),
          places=3,
      )
  
      eval_subdir = os.path.join(self.test_subdirectory, "eval")
      ensemble_eval_subdir = os.path.join(ensemble_subdir, "eval")
      for subdir in [ensemble_eval_subdir, eval_subdir]:
          self.assertEqual(
              [b"| dnn |"],
              tu.check_eventfile_for_keyword(
                  "architecture/adanet/ensembles/0", subdir
              ),
          )
          if subdir == eval_subdir:
              self.assertAlmostEqual(
                  want_loss,
                  tu.check_eventfile_for_keyword("loss", subdir),
                  places=1,
              )
          # TODO: Check why some eval metrics are zero on TPU.
          self.assertAlmostEqual(
              0.0 if use_tpu else want_eval_summary_loss,
              tu.check_eventfile_for_keyword("average_loss", subdir),
              places=1,
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import itertools
  import json
  import os
  from absl.testing import parameterized
  from adanet import tf_compat
  from adanet.core import testing_utils
  from adanet.core.tpu_estimator import TPUEstimator
  from adanet.subnetwork import Builder
  from adanet.subnetwork import Report
  from adanet.subnetwork import SimpleGenerator
  from adanet.subnetwork import Subnetwork
  import numpy
  import tensorflow.compat.v2
  from tensorflow_estimator.python.estimator.head import regression_head
label_desc: Use image node_id 4 for calling the TPUEstimatorTest obj's underlying member method code with example usage: obj.test_tpu_estimator_summaries(use_tpu, want_loss, want_adanet_loss, want_eval_summary_loss, want_predictions) and returns: dict