text: """Test AdaNet estimator single graph implementation for TF 2.
  
  Copyright 2019 The AdaNet Authors. All Rights Reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      https://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os
  
  from absl import logging
  from adanet import tf_compat
  from adanet.core import testing_utils as tu
  from adanet.core.estimator import Estimator
  from adanet.core.report_materializer import ReportMaterializer
  from adanet.subnetwork import Builder
  from adanet.subnetwork import SimpleGenerator
  from adanet.subnetwork import Subnetwork
  import tensorflow.compat.v2 as tf
  
  from tensorflow_estimator.python.estimator.head import regression_head
  
  logging.set_verbosity(logging.INFO)
  
  XOR_FEATURES = [[1., 0.], [0., 0], [0., 1.], [1., 1.]]
  XOR_LABELS = [[1.], [0.], [1.], [0.]]
  
  
  class _SimpleBuilder(Builder):
    """A simple subnetwork builder that takes feature_columns."""
  
    def __init__(self, name, seed=42):
      self._name = name
      self._seed = seed
  
    @property
    def name(self):
      return self._name
  
    def build_subnetwork(self,
                         features,
                         logits_dimension,
                         training,
                         iteration_step,
                         summary,
                         previous_ensemble=None):
      seed = self._seed
      if previous_ensemble:
        # Increment seed so different iterations don't learn the exact same thing.
        seed += 1
  
      with tf_compat.v1.variable_scope("simple"):
        input_layer = tf_compat.v1.feature_column.input_layer(
            features=features,
            feature_columns=tf.feature_column.numeric_column("x", 2))
        last_layer = input_layer
  
      with tf_compat.v1.variable_scope("logits"):
        logits = tf_compat.v1.layers.dense(
            last_layer,
            logits_dimension,
            kernel_initializer=tf_compat.v1.glorot_uniform_initializer(seed=seed))
  
      summary.scalar("scalar", 3)
      batch_size = features["x"].get_shape().as_list()[0]
      summary.image("image", tf.ones([batch_size, 3, 3, 1]))
      with tf_compat.v1.variable_scope("nested"):
        summary.scalar("scalar", 5)
  
      return Subnetwork(
          last_layer=last_layer,
          logits=logits,
          complexity=1,
          persisted_tensors={},
      )
  
    def build_subnetwork_train_op(self, subnetwork, loss, var_list, labels,
                                  iteration_step, summary, previous_ensemble):
      optimizer = tf_compat.v1.train.GradientDescentOptimizer(learning_rate=.001)
      return optimizer.minimize(loss, var_list=var_list)
  
  
  class EstimatorSummaryWriterTest(tu.AdanetTestCase):
    """Test that Tensorboard summaries get written correctly."""
  
    @tf_compat.skip_for_tf1
    def test_summaries(self):
      """Tests that summaries are written to candidate directory."""
  
      run_config = tf.estimator.RunConfig(
          tf_random_seed=42,
          log_step_count_steps=2,
          save_summary_steps=2,
          model_dir=self.test_subdirectory)
      subnetwork_generator = SimpleGenerator([_SimpleBuilder("dnn")])
      report_materializer = ReportMaterializer(
          input_fn=tu.dummy_input_fn([[1., 1.]], [[0.]]), steps=1)
      estimator = Estimator(
          head=regression_head.RegressionHead(
              loss_reduction=tf_compat.SUM_OVER_BATCH_SIZE),
          subnetwork_generator=subnetwork_generator,
          report_materializer=report_materializer,
          max_iteration_steps=10,
          config=run_config)
      train_input_fn = tu.dummy_input_fn([[1., 0.]], [[1.]])
      estimator.train(input_fn=train_input_fn, max_steps=3)
  
      ensemble_loss = 1.52950
      self.assertAlmostEqual(
          ensemble_loss,
          tu.check_eventfile_for_keyword("loss", self.test_subdirectory),
          places=3)
      self.assertIsNotNone(
          tu.check_eventfile_for_keyword("global_step/sec",
                                         self.test_subdirectory))
      self.assertEqual(
          0.,
          tu.check_eventfile_for_keyword("iteration/adanet/iteration",
                                         self.test_subdirectory))
  
      subnetwork_subdir = os.path.join(self.test_subdirectory,
                                       "subnetwork/t0_dnn")
      self.assertAlmostEqual(
          3.,
          tu.check_eventfile_for_keyword("scalar", subnetwork_subdir),
          places=3)
      self.assertEqual((3, 3, 1),
                       tu.check_eventfile_for_keyword("image", subnetwork_subdir))
      self.assertAlmostEqual(
          5.,
          tu.check_eventfile_for_keyword("nested/scalar", subnetwork_subdir),
          places=3)
  
      ensemble_subdir = os.path.join(
          self.test_subdirectory, "ensemble/t0_dnn_grow_complexity_regularized")
      self.assertAlmostEqual(
          ensemble_loss,
          tu.check_eventfile_for_keyword(
              "adanet_loss/adanet/adanet_weighted_ensemble", ensemble_subdir),
          places=1)
      self.assertAlmostEqual(
          0.,
          tu.check_eventfile_for_keyword(
              "complexity_regularization/adanet/adanet_weighted_ensemble",
              ensemble_subdir),
          places=3)
      self.assertAlmostEqual(
          1.,
          tu.check_eventfile_for_keyword(
              "mixture_weight_norms/adanet/"
              "adanet_weighted_ensemble/subnetwork_0", ensemble_subdir),
          places=3)
  
  
  if __name__ == "__main__":
    tf.enable_v2_behavior()
    tf.test.main()
  