message: Add an internal way to turn off batch-size checking on every call to enqueue(). These checks turn out to add a huge amount of time to graph generation:
diff: @@ -341,6 +341,7 @@ class TPUEmbedding(tracking.AutoTrackable):
  self._hosts = get_list_of_hosts(self._strategy)
  self._built = False
  + self._verify_batch_size_on_enqueue = True
  def build(self, per_replica_batch_size: Optional[int] = None):
  """Create the underlying variables and initializes the TPU for embeddings.
  @@ -1232,6 +1233,13 @@ class TPUEmbedding(tracking.AutoTrackable):
  in_tpu_context = self._raise_error_for_incorrect_control_flow_context()
  + if not self._verify_batch_size_on_enqueue:
  + if not self._batch_size or not self._built:
  + raise ValueError(
  + "Configured not to check batch size on each enqueue() call; please "
  + "ensure build() was called with global batch size to initialize "
  + "the TPU for embeddings.")
  + else:
  # Should we also get batch_size from weights if they exist?
  # Since features is assumed to be batched at the per replica batch size
  # the returned batch size here is per replica an not global.
  @@ -1243,8 +1251,8 @@ class TPUEmbedding(tracking.AutoTrackable):
  if batch_size is not None:
  self._maybe_build(batch_size)
  if self._batch_size != batch_size:
  - raise ValueError("Multiple calls to enqueue with different batch sizes "
  - "{} and {}.".format(self._batch_size,
  + raise ValueError("Multiple calls to enqueue with different batch "
  + "sizes {} and {}.".format(self._batch_size,
  batch_size))
  nest.assert_same_structure(self._feature_config, features)
  