title: Colab TPU: TensorFlow '2.0.0-beta0' LinearClassifier .train Bug
body: <p>Attempting to get LinearClassifier running with Colab TPU.
  <a href="https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/estimator/LinearClassifier" rel="nofollow noreferrer">https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/estimator/LinearClassifier</a></p>
  
  <p>TPUStrategy is supported in TensorFlow 2.0 Estimator(LinearClassifier)
  <a href="https://www.tensorflow.org/beta/guide/distribute_strategy#whats_supported_now_2" rel="nofollow noreferrer">https://www.tensorflow.org/beta/guide/distribute_strategy#whats_supported_now_2</a></p>
  
  <p>LinearClassifier works as expected without the tpu_strategy.
  <a href="https://www.tensorflow.org/beta/guide/distribute_strategy#tpustrategy" rel="nofollow noreferrer">https://www.tensorflow.org/beta/guide/distribute_strategy#tpustrategy</a></p>
  
  <p>When adding the tpu_strategy as the config for LinearClassifier i'm getting the following error:</p>
  
  <p><b>InvalidArgumentError: No OpKernel was registered to support Op 'TPUReplicatedInput'</b> used by {{node input0}}with these attrs: [T=DT_DOUBLE, N=8]
  <b>Registered devices: [CPU, XLA_CPU]</b>
  Registered kernels:
  
  [[input0]]
  <a href="https://www.tensorflow.org/beta/guide/distribute_strategy#using_tfdistributestrategy_with_estimator" rel="nofollow noreferrer">https://www.tensorflow.org/beta/guide/distribute_strategy#using_tfdistributestrategy_with_estimator</a></p>
  
  <p>Have been fighting this for a few days now, what is wrong here?</p>
  
  <pre class="lang-py prettyprint-override"><code>!pip install tensorflow==2.0.0-beta0
  import tensorflow.feature_column as fc
  import tensorflow as tf
  import os
  
  print(tf.__version__)
  
  TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=TPU_WORKER)
  tf.config.experimental_connect_to_host(cluster_resolver.master())
  tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
  tpu_strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)
  config = tf.estimator.RunConfig(train_distribute=tpu_strategy, eval_distribute=tpu_strategy)
  
  batch_size = 1
  def make_input_fn(X, y):
    def input_fn():
      dataset = tf.data.Dataset.from_tensor_slices((dict(X), y)).batch(batch_size)
      return dataset
    return input_fn
  
  input_fn = make_input_fn(estimator_train_attributes_dictionary,labels_train)
  linear_est = tf.estimator.LinearClassifier(feature_columns=attibute_columns,config=config)
  linear_est.train(input_fn=input_fn)
  </code></pre>
  
  <p>full exception in Colab:</p>
  
  <pre class="lang-py prettyprint-override"><code>W0618 18:08:10.280844 140506166175616 estimator.py:1811] Using temporary folder as model directory: /tmp/tmp2xc1fixj
  2.0.0-beta0
  W0618 18:09:00.986362 140506166175616 tpu.py:218] 3 unsupported operations found: 
    ScalarSummary (bias)
    ScalarSummary (fraction_of_zero_weights)
    ScalarSummary (loss)
  W0618 18:09:43.578035 140506166175616 tpu_strategy_util.py:57] TPU system %s has already been initialized. Reinitializing the TPU can cause previously created variables on TPU to be lost.
  ---------------------------------------------------------------------------
  InvalidArgumentError                      Traceback (most recent call last)
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _do_call(self, fn, *args)
     1355     try:
  -&gt; 1356       return fn(*args)
     1357     except errors.OpError as e:
  
  20 frames
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _run_fn(feed_dict, fetch_list, target_list, options, run_metadata)
     1338       # Ensure any changes to the graph are reflected in the runtime.
  -&gt; 1339       self._extend_graph()
     1340       return self._call_tf_sessionrun(
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _extend_graph(self)
     1373     with self._graph._session_run_lock():  # pylint: disable=protected-access
  -&gt; 1374       tf_session.ExtendSession(self._session)
     1375 
  
  InvalidArgumentError: No OpKernel was registered to support Op 'TPUReplicatedInput' used by {{node input0}}with these attrs: [T=DT_DOUBLE, N=8]
  Registered devices: [CPU, XLA_CPU]
  Registered kernels:
    &lt;no registered kernels&gt;
  
       [[input0]]
  
  During handling of the above exception, another exception occurred:
  
  InvalidArgumentError                      Traceback (most recent call last)
  &lt;ipython-input-23-66caf93d8677&gt; in &lt;module&gt;()
       25 
       26 linear_est = tf.estimator.LinearClassifier(feature_columns=attibute_columns,config=config)#feature_columns=featureNames,,config=config
  ---&gt; 27 linear_est.train(input_fn=input_fn)#,max_steps=100
       28 
       29 #train_spec = tf.estimator.TrainSpec(input_fn=input_fn, max_steps=1000)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py in train(self, input_fn, hooks, steps, max_steps, saving_listeners)
      365 
      366       saving_listeners = _check_listeners_type(saving_listeners)
  --&gt; 367       loss = self._train_model(input_fn, hooks, saving_listeners)
      368       logging.info('Loss for final step: %s.', loss)
      369       return self
  
  /usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py in _train_model(self, input_fn, hooks, saving_listeners)
     1154   def _train_model(self, input_fn, hooks, saving_listeners):
     1155     if self._train_distribution:
  -&gt; 1156       return self._train_model_distributed(input_fn, hooks, saving_listeners)
     1157     else:
     1158       return self._train_model_default(input_fn, hooks, saving_listeners)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py in _train_model_distributed(self, input_fn, hooks, saving_listeners)
     1217       self._config._train_distribute.configure(self._config.session_config)
     1218       return self._actual_train_model_distributed(
  -&gt; 1219           self._config._train_distribute, input_fn, hooks, saving_listeners)
     1220     # pylint: enable=protected-access
     1221 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py in _actual_train_model_distributed(self, strategy, input_fn, hooks, saving_listeners)
     1327         return self._train_with_estimator_spec(estimator_spec, worker_hooks,
     1328                                                hooks, global_step_tensor,
  -&gt; 1329                                                saving_listeners)
     1330 
     1331   def _train_with_estimator_spec_distributed(self, estimator_spec, worker_hooks,
  
  /usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py in _train_with_estimator_spec(self, estimator_spec, worker_hooks, hooks, global_step_tensor, saving_listeners)
     1478         save_summaries_steps=save_summary_steps,
     1479         config=self._session_config,
  -&gt; 1480         log_step_count_steps=log_step_count_steps) as mon_sess:
     1481       loss = None
     1482       any_step_done = False
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in MonitoredTrainingSession(master, is_chief, checkpoint_dir, scaffold, hooks, chief_only_hooks, save_checkpoint_secs, save_summaries_steps, save_summaries_secs, config, stop_grace_period_secs, log_step_count_steps, max_wait_secs, save_checkpoint_steps, summary_dir)
      582       session_creator=session_creator,
      583       hooks=all_hooks,
  --&gt; 584       stop_grace_period_secs=stop_grace_period_secs)
      585 
      586 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in __init__(self, session_creator, hooks, stop_grace_period_secs)
     1005         hooks,
     1006         should_recover=True,
  -&gt; 1007         stop_grace_period_secs=stop_grace_period_secs)
     1008 
     1009 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in __init__(self, session_creator, hooks, should_recover, stop_grace_period_secs)
      723         stop_grace_period_secs=stop_grace_period_secs)
      724     if should_recover:
  --&gt; 725       self._sess = _RecoverableSession(self._coordinated_creator)
      726     else:
      727       self._sess = self._coordinated_creator.create_session()
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in __init__(self, sess_creator)
     1198     """
     1199     self._sess_creator = sess_creator
  -&gt; 1200     _WrappedSession.__init__(self, self._create_session())
     1201 
     1202   def _create_session(self):
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in _create_session(self)
     1203     while True:
     1204       try:
  -&gt; 1205         return self._sess_creator.create_session()
     1206       except _PREEMPTION_ERRORS as e:
     1207         logging.info(
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in create_session(self)
      869       """Creates a coordinated session."""
      870       # Keep the tf_sess for unit testing.
  --&gt; 871       self.tf_sess = self._session_creator.create_session()
      872       # We don't want coordinator to suppress any exception.
      873       self.coord = coordinator.Coordinator(clean_stop_exception_types=[])
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/monitored_session.py in create_session(self)
      645         init_op=self._scaffold.init_op,
      646         init_feed_dict=self._scaffold.init_feed_dict,
  --&gt; 647         init_fn=self._scaffold.init_fn)
      648 
      649 
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/session_manager.py in prepare_session(self, master, init_op, saver, checkpoint_dir, checkpoint_filename_with_path, wait_for_checkpoint, max_wait_secs, config, init_feed_dict, init_fn)
      294                            "init_fn or local_init_op was given")
      295       if init_op is not None:
  --&gt; 296         sess.run(init_op, feed_dict=init_feed_dict)
      297       if init_fn:
      298         init_fn(sess)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in run(self, fetches, feed_dict, options, run_metadata)
      948     try:
      949       result = self._run(None, fetches, feed_dict, options_ptr,
  --&gt; 950                          run_metadata_ptr)
      951       if run_metadata:
      952         proto_data = tf_session.TF_GetBuffer(run_metadata_ptr)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _run(self, handle, fetches, feed_dict, options, run_metadata)
     1171     if final_fetches or final_targets or (handle and feed_dict_tensor):
     1172       results = self._do_run(handle, final_targets, final_fetches,
  -&gt; 1173                              feed_dict_tensor, options, run_metadata)
     1174     else:
     1175       results = []
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _do_run(self, handle, target_list, fetch_list, feed_dict, options, run_metadata)
     1348     if handle is None:
     1349       return self._do_call(_run_fn, feeds, fetches, targets, options,
  -&gt; 1350                            run_metadata)
     1351     else:
     1352       return self._do_call(_prun_fn, handle, feeds, fetches)
  
  /usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py in _do_call(self, fn, *args)
     1368           pass
     1369       message = error_interpolation.interpolate(message, self._graph)
  -&gt; 1370       raise type(e)(node_def, op, message)
     1371 
     1372   def _extend_graph(self):
  
  
  InvalidArgumentError: No OpKernel was registered to support Op 'TPUReplicatedInput' used by node input0 (defined at &lt;ipython-input-23-66caf93d8677&gt;:27) with these attrs: [T=DT_DOUBLE, N=8]
  Registered devices: [CPU, XLA_CPU]
  Registered kernels:
    &lt;no registered kernels&gt;
  
       [[input0]]
  </code></pre>