message: [TPU] Add extra type checking in Python around the infeed operation, since it is a common mistake to make.
diff: @@ -21,6 +21,7 @@ from __future__ import print_function
  import platform
  +from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  if platform.system() != "Windows":
  @@ -40,6 +41,63 @@ if platform.system() != "Windows":
  del op # Unused
  # The gradient of a cross replica sum is also a cross-replica sum.
  return gen_tpu_ops.cross_replica_sum(grad)
  +
  + # This extra type checking exists to give a more helpful error message in
  + # the common case that uint8 and int64 values are infed. Remove when both
  + # types are supported.
  +
  + _SUPPORTED_INFEED_DTYPES = set([
  + dtypes.int32, dtypes.bfloat16, dtypes.float32
  + ])
  +
  + def infeed_dequeue(dtype, shape, name=None):
  + """A placeholder op for a value that will be fed into the computation.
  +
  + Args:
  + dtype: A `tf.DType`. The type of elements in the tensor.
  + shape: A `tf.TensorShape` or list of `ints`. The shape of the tensor.
  + name: A name for the operation (optional).
  +
  + Returns:
  + A `Tensor` of type `dtype`.
  + A tensor that will be provided using the infeed mechanism.
  +
  + Raises:
  + TypeError: If 'dtype` is not a supported infeed type.
  + """
  + if dtype not in _SUPPORTED_INFEED_DTYPES:
  + raise TypeError(
  + "{} is not a supported TPU infeed type. Supported types are: "
  + "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES)))
  +
  + return gen_tpu_ops.infeed_dequeue(dtype, shape, name=name)
  +
  + # pylint: disable=redefined-outer-name
  + def infeed_dequeue_tuple(dtypes, shapes, name=None):
  + """A placeholder op for values fed into the TPU simultaneously as a tuple.
  +
  + Args:
  + dtypes: A list of `tf.DType`s that has length `>= 1`.
  + The element types of each element in `outputs`.
  + shapes: A list of shapes (each a `tf.TensorShape` or list of `ints`).
  + The shapes of each tensor in `outputs`.
  + name: A name for the operation (optional).
  +
  + Returns:
  + A list of `Tensor` objects of type `dtypes`.
  + A list of tensors that will be provided using the infeed mechanism.
  +
  + Raises:
  + TypeError: If a type in 'dtypes` is not a supported infeed type.
  + """
  + for dtype in dtypes:
  + if dtype not in _SUPPORTED_INFEED_DTYPES:
  + raise TypeError(
  + "{} is not a supported TPU infeed type. Supported types are: "
  + "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES)))
  + return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name)
  + # pylint: enable=redefined-outer-name
  +
  else:
  # We have already built the appropriate libraries into the binary via CMake
  # if we have built contrib, so we don't need this
  