# Copyright 2017 The TensorFlow 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
#
#     http://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.
# =============================================================================

"""Operations for TPUs."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import platform

from tensorflow.python.framework import ops

if platform.system() != "Windows":
  # pylint: disable=wildcard-import,unused-import,g-import-not-at-top
  from tensorflow.contrib.tpu.ops import gen_tpu_ops
  from tensorflow.contrib.tpu.ops.gen_tpu_ops import *

  from tensorflow.contrib.util import loader
  from tensorflow.python.platform import resource_loader
  # pylint: enable=wildcard-import,unused-import,g-import-not-at-top

  _tpu_ops = loader.load_op_library(
      resource_loader.get_path_to_datafile("_tpu_ops.so"))

  @ops.RegisterGradient("CrossReplicaSum")
  def _cross_replica_sum_grad(op, grad):
    del op  # Unused
    # The gradient of a cross replica sum is also a cross-replica sum.
    return gen_tpu_ops.cross_replica_sum(grad)
else:
  # We have already built the appropriate libraries into the binary via CMake
  # if we have built contrib, so we don't need this
  pass

@@ -837,16 +837,63 @@
 atform%0A%0A
+from tensorflow.python.framework import dtypes%0A
 from ten
@@ -1635,16 +1635,2127 @@
 m(grad)%0A
+%0A  # This extra type checking exists to give a more helpful error message in%0A  # the common case that uint8 and int64 values are infed. Remove when both%0A  # types are supported.%0A%0A  _SUPPORTED_INFEED_DTYPES = set(%5B%0A      dtypes.int32, dtypes.bfloat16, dtypes.float32%0A  %5D)%0A%0A  def infeed_dequeue(dtype, shape, name=None):%0A    %22%22%22A placeholder op for a value that will be fed into the computation.%0A%0A    Args:%0A      dtype: A %60tf.DType%60. The type of elements in the tensor.%0A      shape: A %60tf.TensorShape%60 or list of %60ints%60. The shape of the tensor.%0A      name: A name for the operation (optional).%0A%0A    Returns:%0A      A %60Tensor%60 of type %60dtype%60.%0A      A tensor that will be provided using the infeed mechanism.%0A%0A    Raises:%0A      TypeError: If 'dtype%60 is not a supported infeed type.%0A    %22%22%22%0A    if dtype not in _SUPPORTED_INFEED_DTYPES:%0A      raise TypeError(%0A          %22%7B%7D is not a supported TPU infeed type. Supported types are: %22%0A          %22%7B%7D%22.format(dtype, list(_SUPPORTED_INFEED_DTYPES)))%0A%0A    return gen_tpu_ops.infeed_dequeue(dtype, shape, name=name)%0A%0A  # pylint: disable=redefined-outer-name%0A  def infeed_dequeue_tuple(dtypes, shapes, name=None):%0A    %22%22%22A placeholder op for values fed into the TPU simultaneously as a tuple.%0A%0A    Args:%0A      dtypes: A list of %60tf.DType%60s that has length %60%3E= 1%60.%0A        The element types of each element in %60outputs%60.%0A      shapes: A list of shapes (each a %60tf.TensorShape%60 or list of %60ints%60).%0A        The shapes of each tensor in %60outputs%60.%0A      name: A name for the operation (optional).%0A%0A    Returns:%0A      A list of %60Tensor%60 objects of type %60dtypes%60.%0A      A list of tensors that will be provided using the infeed mechanism.%0A%0A    Raises:%0A      TypeError: If a type in 'dtypes%60 is not a supported infeed type.%0A    %22%22%22%0A    for dtype in dtypes:%0A      if dtype not in _SUPPORTED_INFEED_DTYPES:%0A        raise TypeError(%0A            %22%7B%7D is not a supported TPU infeed type. Supported types are: %22%0A            %22%7B%7D%22.format(dtype, list(_SUPPORTED_INFEED_DTYPES)))%0A    return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name)%0A  # pylint: enable=redefined-outer-name%0A%0A
 else:%0A  

Add extra type checking in Python around the infeed operation, since it is a common mistake to make.