Copy TPUEstimator to Estimator.

Expose TPUEstimator symbols under tf.compat.v1.tpu.

PiperOrigin-RevId: 243871401

tensorflow/python/BUILD

@@ -198,6 +198,7 @@ py_library(
         "//tensorflow/python/saved_model",
         "//tensorflow/python/tools:component_api_helper",
         "//tensorflow/python/tools/api/generator:create_python_api",
+        "//tensorflow/python/tpu:tpu_noestimator",
         "//third_party/py/numpy",
     ],
 )

tensorflow/python/tpu/BUILD

@@ -132,6 +132,20 @@ py_library(
     ],
 )
 
+py_library(
+    name = "tpu_noestimator",
+    srcs = [
+        "__init__.py",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":feature_column",
+        ":preempted_hook_py",
+        ":tpu_embedding",
+        ":tpu_lib",
+    ],
+)
+
 py_library(
     name = "tpu_lib",
     srcs = [

tensorflow/tools/api/golden/v1/tensorflow.estimator.pbtxt

@@ -188,6 +188,10 @@ tf_module {
     name: "inputs"
     mtype: "<type \'module\'>"
   }
+  member {
+    name: "tpu"
+    mtype: "<type \'module\'>"
+  }
   member_method {
     name: "add_metrics"
     argspec: "args=[\'estimator\', \'metric_fn\'], varargs=None, keywords=None, defaults=None"

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.-input-pipeline-config.pbtxt

@@ -0,0 +1,28 @@
+path: "tensorflow.estimator.tpu.InputPipelineConfig"
+tf_class {
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_config.InputPipelineConfig\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "BROADCAST"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "PER_HOST_V1"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "PER_HOST_V2"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "PER_SHARD_V1"
+    mtype: "<type \'int\'>"
+  }
+  member {
+    name: "SLICED"
+    mtype: "<type \'int\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+}

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.-run-config.pbtxt

@@ -0,0 +1,118 @@
+path: "tensorflow.estimator.tpu.RunConfig"
+tf_class {
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_config.RunConfig\'>"
+  is_instance: "<class \'tensorflow_estimator.python.estimator.run_config.RunConfig\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "cluster"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "cluster_spec"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "device_fn"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "eval_distribute"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "evaluation_master"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "experimental_max_worker_delay_secs"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "global_id_in_cluster"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "is_chief"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "keep_checkpoint_every_n_hours"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "keep_checkpoint_max"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "log_step_count_steps"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "master"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "model_dir"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "num_ps_replicas"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "num_worker_replicas"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "protocol"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "save_checkpoints_secs"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "save_checkpoints_steps"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "save_summary_steps"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "service"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "session_config"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "task_id"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "task_type"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "tf_random_seed"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "tpu_config"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "train_distribute"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'tpu_config\', \'evaluation_master\', \'master\', \'cluster\'], varargs=None, keywords=kwargs, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+  }
+  member_method {
+    name: "replace"
+    argspec: "args=[\'self\'], varargs=None, keywords=kwargs, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.-t-p-u-config.pbtxt

@@ -0,0 +1,47 @@
+path: "tensorflow.estimator.tpu.TPUConfig"
+tf_class {
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_config.TPUConfig\'>"
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_config.TPUConfig\'>"
+  is_instance: "<type \'tuple\'>"
+  member {
+    name: "eval_training_input_configuration"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "initial_infeed_sleep_secs"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "input_partition_dims"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "iterations_per_loop"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "num_cores_per_replica"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "num_shards"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "per_host_input_for_training"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "tpu_job_name"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "count"
+  }
+  member_method {
+    name: "index"
+  }
+}

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.-t-p-u-estimator-spec.pbtxt

@@ -0,0 +1,64 @@
+path: "tensorflow.estimator.tpu.TPUEstimatorSpec"
+tf_class {
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_estimator.TPUEstimatorSpec\'>"
+  is_instance: "<class \'tensorflow_estimator.python.estimator.model_fn._TPUEstimatorSpec\'>"
+  is_instance: "<class \'tensorflow_estimator.python.estimator.model_fn.TPUEstimatorSpec\'>"
+  is_instance: "<type \'tuple\'>"
+  member {
+    name: "eval_metrics"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "evaluation_hooks"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "export_outputs"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "host_call"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "loss"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "mode"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "prediction_hooks"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "predictions"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "scaffold_fn"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "train_op"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "training_hooks"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "as_estimator_spec"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "count"
+  }
+  member_method {
+    name: "index"
+  }
+}

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.-t-p-u-estimator.pbtxt

@@ -0,0 +1,66 @@
+path: "tensorflow.estimator.tpu.TPUEstimator"
+tf_class {
+  is_instance: "<class \'tensorflow_estimator.python.estimator.tpu.tpu_estimator.TPUEstimator\'>"
+  is_instance: "<class \'tensorflow_estimator.python.estimator.estimator.Estimator\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "config"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "model_dir"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "model_fn"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "params"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+    argspec: "args=[\'self\', \'model_fn\', \'model_dir\', \'config\', \'params\', \'use_tpu\', \'train_batch_size\', \'eval_batch_size\', \'predict_batch_size\', \'batch_axis\', \'eval_on_tpu\', \'export_to_tpu\', \'export_to_cpu\', \'warm_start_from\', \'experimental_export_device_assignment\', \'experimental_embedding_config_spec\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\', \'True\', \'None\', \'None\', \'None\', \'None\', \'True\', \'True\', \'True\', \'None\', \'False\', \'None\'], "
+  }
+  member_method {
+    name: "eval_dir"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+  }
+  member_method {
+    name: "evaluate"
+    argspec: "args=[\'self\', \'input_fn\', \'steps\', \'hooks\', \'checkpoint_path\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+  }
+  member_method {
+    name: "experimental_export_all_saved_models"
+    argspec: "args=[\'self\', \'export_dir_base\', \'input_receiver_fn_map\', \'assets_extra\', \'as_text\', \'checkpoint_path\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\'], "
+  }
+  member_method {
+    name: "export_saved_model"
+    argspec: "args=[\'self\', \'export_dir_base\', \'serving_input_receiver_fn\', \'assets_extra\', \'as_text\', \'checkpoint_path\', \'experimental_mode\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\', \'infer\'], "
+  }
+  member_method {
+    name: "export_savedmodel"
+    argspec: "args=[\'self\', \'export_dir_base\', \'serving_input_receiver_fn\', \'assets_extra\', \'as_text\', \'checkpoint_path\', \'strip_default_attrs\'], varargs=None, keywords=None, defaults=[\'None\', \'False\', \'None\', \'False\'], "
+  }
+  member_method {
+    name: "get_variable_names"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "get_variable_value"
+    argspec: "args=[\'self\', \'name\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "latest_checkpoint"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "predict"
+    argspec: "args=[\'self\', \'input_fn\', \'predict_keys\', \'hooks\', \'checkpoint_path\', \'yield_single_examples\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'True\'], "
+  }
+  member_method {
+    name: "train"
+    argspec: "args=[\'self\', \'input_fn\', \'hooks\', \'steps\', \'max_steps\', \'saving_listeners\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\', \'None\'], "
+  }
+}

tensorflow/tools/api/golden/v1/tensorflow.estimator.tpu.pbtxt

@@ -0,0 +1,23 @@
+path: "tensorflow.estimator.tpu"
+tf_module {
+  member {
+    name: "InputPipelineConfig"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "RunConfig"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "TPUConfig"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "TPUEstimator"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "TPUEstimatorSpec"
+    mtype: "<type \'type\'>"
+  }
+}

tensorflow/tools/compatibility/renames_v2.py

@@ -345,6 +345,16 @@ renames = {
         'tf.compat.v1.estimator.inputs.numpy_input_fn',
     'tf.estimator.inputs.pandas_input_fn':
         'tf.compat.v1.estimator.inputs.pandas_input_fn',
+    'tf.estimator.tpu.InputPipelineConfig':
+        'tf.compat.v1.estimator.tpu.InputPipelineConfig',
+    'tf.estimator.tpu.RunConfig':
+        'tf.compat.v1.estimator.tpu.RunConfig',
+    'tf.estimator.tpu.TPUConfig':
+        'tf.compat.v1.estimator.tpu.TPUConfig',
+    'tf.estimator.tpu.TPUEstimator':
+        'tf.compat.v1.estimator.tpu.TPUEstimator',
+    'tf.estimator.tpu.TPUEstimatorSpec':
+        'tf.compat.v1.estimator.tpu.TPUEstimatorSpec',
     'tf.expm1':
         'tf.math.expm1',
     'tf.fake_quant_with_min_max_args':