class_docstr: An :class:`adanet.Estimator` capable of training and evaluating on TPU.
  
  Unless :code:`use_tpu=False`, training will run on TPU. However, certain parts
  of the AdaNet training loop, such as report materialization and best candidate
  selection, will still occurr on CPU. Furthermore, if using TPUEmbedding (i.e.
  :code:`embedding_config_spec` is supplied), inference will also occurr on CPU.
  
  TODO: Provide the missing functionality detailed below.
  N.B: Embeddings using the TPUEmbedding (i.e. :code:`embedding_config_spec`
  is provided) only support :code:`shared_embedding_columns` when running for
  multiple AdaNet iterations. Using regular :code:`embedding_columns` will cause
  iterations 2..n to fail because of mismatched embedding scopes.
  
  Args:
    head: See :class:`adanet.Estimator`.
    subnetwork_generator: See :class:`adanet.Estimator`.
    max_iteration_steps: See :class:`adanet.Estimator`.
    ensemblers: See :class:`adanet.Estimator`.
    ensemble_strategies: See :class:`adanet.Estimator`.
    evaluator: See :class:`adanet.Estimator`.
    report_materializer: See :class:`adanet.Estimator`.
    metric_fn: See :class:`adanet.Estimator`.
    force_grow: See :class:`adanet.Estimator`.
    replicate_ensemble_in_training: See :class:`adanet.Estimator`.
    adanet_loss_decay: See :class:`adanet.Estimator`.
    report_dir: See :class:`adanet.Estimator`.
    config: See :class:`adanet.Estimator`.
    use_tpu: Boolean to enable training on TPU. Defaults to :code:`True` and is
      only provided to allow debugging models on CPU/GPU. Use
      :class:`adanet.Estimator` instead if you do not plan to run on TPU.
    eval_on_tpu: Boolean to enable evaluating on TPU. Defaults to :code:`True`.
      Ignored if :code:`use_tpu=False`.
    export_to_tpu: See :class:`tf.compat.v1.estimator.tpu.TPUEstimator`.
    train_batch_size: See :class:`tf.compat.v1.estimator.tpu.TPUEstimator`.
      Defaults to 0 if `None`.
    eval_batch_size: See :class:`tf.compat.v1.estimator.tpu.TPUEstimator`.
      Defaults to train_batch_size if `None`.
    predict_batch_size: See :class:`tf.compat.v1.estimator.tpu.TPUEstimator`.
      Defaults to eval_batch_size if `None`.
    embedding_config_spec: See :class:`tf.compat.v1.estimator.tpu.TPUEstimator`.
      If supplied, :code:`predict` will be called on CPU and no TPU compatible
        :code:`SavedModel` will be exported.
    debug: See :class:`adanet.Estimator`.
    enable_ensemble_summaries: See :class:`adanet.Estimator`.
    enable_subnetwork_summaries: See :class:`adanet.Estimator`.
    export_subnetwork_logits: Whether to include subnetwork logits in exports.
    export_subnetwork_last_layer: Whether to include subnetwork last layer in
      exports.
    global_step_combiner_fn: See :class:`adanet.Estimator`.
    max_iterations: See :class:`adanet.Estimator`.
    replay_config: See :class:`adanet.Estimator`.
    add_predict_batch_config: If True, supplies a default
      `tpu_estimator.BatchConfig` when calling
      `tpu_estimator.model_fn_inference_on_tpu`, otherwise supplies None.
    **kwargs: Extra keyword args passed to the parent.
code: def __init__(
      self,
      head,
      subnetwork_generator,
      max_iteration_steps,
      ensemblers=None,
      ensemble_strategies=None,
      evaluator=None,
      report_materializer=None,
      metric_fn=None,
      force_grow=False,
      replicate_ensemble_in_training=False,
      adanet_loss_decay=0.9,
      model_dir=None,
      report_dir=None,
      config=None,
      use_tpu=True,
      eval_on_tpu=True,
      export_to_tpu=True,
      train_batch_size=None,
      eval_batch_size=None,
      predict_batch_size=None,
      embedding_config_spec=None,
      debug=False,
      enable_ensemble_summaries=True,
      enable_subnetwork_summaries=True,
      export_subnetwork_logits=False,
      export_subnetwork_last_layer=True,
      global_step_combiner_fn=tf.math.reduce_mean,
      max_iterations=None,
      replay_config=None,
      add_predict_batch_config=True,
      **kwargs,
  ):
      self._use_tpu = use_tpu
      if not self._use_tpu:
          logging.warning(
              "This adanet.TPUEstimator is meant to be used for running on TPU. "
              "If you want to run on CPU/GPU, use adanet.Estimator instead."
          )
      # TPUEstimator modifies config under the hood. We keep track of it here so
      # we can use it from _create_temp_run_config.
      self._original_config = (
          config or tf_compat.v1.estimator.tpu.RunConfig()
      )
      self._eval_on_tpu = eval_on_tpu if self._use_tpu else False
      self._export_to_tpu = export_to_tpu
      self._train_batch_size = train_batch_size or 0
      self._eval_batch_size = eval_batch_size or train_batch_size or 0
      self._predict_batch_size = (
          predict_batch_size or eval_batch_size or train_batch_size or 0
      )
      self._embedding_config_spec = embedding_config_spec
      self._add_predict_batch_config = add_predict_batch_config
      if self._embedding_config_spec:
          logging.warning(
              "TPU does not support inference with TPUEmbedding. Force setting "
              "`export_to_tpu=False` so no TPU SavedModel will be exported."
          )
          self._export_to_tpu = False
  
      from tensorflow_estimator.python.estimator.tpu import (
          tpu_estimator,
      )  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  
      super(TPUEstimator, self).__init__(
          head=head,
          subnetwork_generator=subnetwork_generator,
          max_iteration_steps=max_iteration_steps,
          ensemblers=ensemblers,
          ensemble_strategies=ensemble_strategies,
          evaluator=evaluator,
          report_materializer=report_materializer,
          metric_fn=metric_fn,
          force_grow=force_grow,
          replicate_ensemble_in_training=replicate_ensemble_in_training,
          adanet_loss_decay=adanet_loss_decay,
          model_dir=model_dir,
          report_dir=report_dir,
          config=self._original_config,
          use_tpu=self._use_tpu,
          eval_on_tpu=self._eval_on_tpu,
          export_to_tpu=self._export_to_tpu,
          export_saved_model_api_version=(
              tpu_estimator.ExportSavedModelApiVersion.V2
          ),
          train_batch_size=self._train_batch_size,
          eval_batch_size=self._eval_batch_size,
          predict_batch_size=self._predict_batch_size,
          embedding_config_spec=self._embedding_config_spec,
          debug=debug,
          enable_ensemble_summaries=enable_ensemble_summaries,
          enable_subnetwork_summaries=enable_subnetwork_summaries,
          export_subnetwork_logits=export_subnetwork_logits,
          export_subnetwork_last_layer=export_subnetwork_last_layer,
          global_step_combiner_fn=global_step_combiner_fn,
          max_iterations=max_iterations,
          replay_config=replay_config,
          **kwargs,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import contextlib
  import functools
  from absl import logging
  from adanet import tf_compat
  from adanet.core.estimator import Estimator
  import tensorflow.compat.v2
label_desc: Use image node_id 1 to create a new TPUEstimator object from inherited base classes: Estimator, tf.compat.v1.estimator.tpu with example: obj = TPUEstimator(head, subnetwork_generator, max_iteration_steps, ensemblers, ensemble_strategies, evaluator, report_materializer, metric_fn, force_grow, replicate_ensemble_in_training, adanet_loss_decay, model_dir, report_dir, config, use_tpu, eval_on_tpu, export_to_tpu, train_batch_size, eval_batch_size, predict_batch_size, embedding_config_spec, debug, enable_ensemble_summaries, enable_subnetwork_summaries, export_subnetwork_logits, export_subnetwork_last_layer, global_step_combiner_fn, max_iterations, replay_config, add_predict_batch_config)