message: Handles the case when number of hosts is more than the number of ids in tpu_embedding.py.
diff: @@ -34,6 +34,7 @@ from tensorflow.python.ops import init_ops
  from tensorflow.python.ops import partitioned_variables
  from tensorflow.python.ops import state_ops
  from tensorflow.python.ops import variable_scope
  +from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.tpu import tpu_system_metadata as tpu_system_metadata_lib
  from tensorflow.python.tpu.ops import tpu_ops
  @@ -606,7 +607,10 @@ class TPUEmbedding(object):
  table_descriptor.name = table
  table_config = self._table_to_config_dict[table]
  - table_descriptor.vocabulary_size = table_config.vocabulary_size
  + # For small tables, we pad to the number of hosts so that at least one
  + # id will be assigned to each host.
  + table_descriptor.vocabulary_size = max(table_config.vocabulary_size,
  + len(self.hosts))
  table_descriptor.dimension = table_config.dimension
  table_descriptor.num_features = self._table_to_num_features_dict[table]
  @@ -1278,14 +1282,19 @@ def _create_device_fn(hosts):
  def device_fn(op):
  """Returns the `device` for `op`."""
  part_match = re.match(r'.*/part_(\d+)(/|$)', op.name)
  + dummy_match = re.match(r'.*dummy_(\d+).*', op.name)
  + if not part_match and not dummy_match:
  + raise RuntimeError(
  + 'Internal Error: Expected {} to contain /part_* or dummy_*'.format(
  + op.name))
  if part_match:
  idx = int(part_match.group(1))
  else:
  - raise RuntimeError('Internal Error: '
  - 'Expected %s to contain /part_*.' % op.name)
  + idx = int(dummy_match.group(1))
  device = hosts[idx]
  + logging.debug('assigning {} to {}.', op, device)
  return device
  return device_fn
  @@ -1298,17 +1307,31 @@ def _create_partitioned_variables(name,
  initializer,
  collections=None): # pylint: disable=redefined-outer-name
  """Creates ParitionedVariables based on `num_hosts` for `table`."""
  - # TODO(shizhiw): automatically place embedding lookup elsewhere?
  - if vocabulary_size < num_hosts:
  - raise ValueError('`vocabulary_size`({}) is smaller than `num_hosts`({}). '
  - 'As TPU embedding is not optimized for small tables, '
  - 'please consider other ways for this embedding lookup.')
  - return list(variable_scope.get_variable(
  + num_slices = min(vocabulary_size, num_hosts)
  +
  + var_list = list(
  + variable_scope.get_variable(
  name,
  shape=(vocabulary_size, embedding_dimension),
  - partitioner=partitioned_variables.fixed_size_partitioner(num_hosts),
  + partitioner=partitioned_variables.fixed_size_partitioner(num_slices),
  dtype=dtypes.float32,
  initializer=initializer,
  collections=collections,
  trainable=False))
  +
  + if vocabulary_size >= num_hosts:
  + return var_list
  +
  + # For padded part, define the dummy variable to be loaded into TPU system.
  + for idx in range(num_hosts - vocabulary_size):
  + var_list.append(
  + variable_scope.get_variable(
  + 'dummy_{}_{}'.format(vocabulary_size + idx, name),
  + shape=(1, embedding_dimension),
  + dtype=dtypes.float32,
  + initializer=initializer,
  + collections=[ops.GraphKeys.LOCAL_VARIABLES],
  + trainable=False))
  +
  + return var_list
  