nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: OpenNMT/OpenNMT-tf
Pull request #512 (2019-10-09 10:31:36)

Fix early stop logic
The improvement should be checked against the steps-1 value only.


Patch:
diff --git a/opennmt/evaluation.py b/opennmt/evaluation.py
index 5dff5acd..34002de9 100644
--- a/opennmt/evaluation.py
+++ b/opennmt/evaluation.py
@@ -312,14 +312,8 @@ def early_stop(metrics, steps, min_improvement=0, higher_is_better=False):
     else:
       return new < ref - min_improvement
 
-  samples = metrics[-steps - 1:]
-
-  # First check if the boundaries show an improvement.
-  if _did_improve(samples[0], samples[-1]):
-    return False
-
-  # If not, only early stop if each successive evaluation did not improve.
-  for metric, next_metric in zip(samples[:-1], samples[1:]):
-    if _did_improve(metric, next_metric):
+  ref_metric = metrics[-steps - 1]
+  for metric in metrics[-steps:]:
+    if _did_improve(ref_metric, metric):
       return False
   return True


Test patch:
diff --git a/opennmt/tests/evaluation_test.py b/opennmt/tests/evaluation_test.py
index 21825eb9..14b1473b 100644
--- a/opennmt/tests/evaluation_test.py
+++ b/opennmt/tests/evaluation_test.py
@@ -144,6 +144,10 @@ class EvaluationTest(tf.test.TestCase):
         evaluation.early_stop([32, 33, 32, 33, 32], 3, min_improvement=2, higher_is_better=False))
     self.assertFalse(
         evaluation.early_stop([32, 35, 32, 33, 32], 3, min_improvement=2, higher_is_better=False))
+    self.assertTrue(
+        evaluation.early_stop(
+            [50.349343, 50.553991, 50.436176, 50.419565, 50.219028, 50.375434],
+            4, min_improvement=0.01, higher_is_better=True))
 
 
 if __name__ == "__main__":
Repository: OpenNMT/OpenNMT-tf
Pull request #522 (2019-10-16 13:32:04)

Do not apply noise module on empty inputs


Patch:
diff --git a/opennmt/data/noise.py b/opennmt/data/noise.py
index d1b378de..76f7acb5 100644
--- a/opennmt/data/noise.py
+++ b/opennmt/data/noise.py
@@ -122,7 +122,11 @@ class Noise(object):
     inputs = words
     if words.shape.ndims == 1:
       inputs = tf.expand_dims(inputs, 1)
-    outputs = self._apply(inputs)
+    num_words = tf.shape(inputs)[0]
+    outputs = tf.cond(
+        tf.math.equal(num_words, 0),
+        true_fn=lambda: inputs,
+        false_fn=lambda: self._apply(inputs))
     if words.shape.ndims == 1:
       outputs = tf.squeeze(outputs, 1)
     return outputs


Test patch:
diff --git a/opennmt/tests/noise_test.py b/opennmt/tests/noise_test.py
index 4c579254..0b4faf9a 100644
--- a/opennmt/tests/noise_test.py
+++ b/opennmt/tests/noise_test.py
@@ -20,14 +20,15 @@ class NoiseTest(tf.test.TestCase):
     self.assertAllEqual(x, y)
 
   @parameterized.expand([
+    [[]],
     [["a", "b", "c", "e", "f"]],
     [[["a", "b", ""], ["c", "e", "f"]]],
   ])
   def testWordDropoutAll(self, words):
-    x = tf.constant(words)
+    x = tf.constant(words, dtype=tf.string)
     y = noise.WordDropout(1)(x)
     y = self.evaluate(y)
-    self.assertEqual(y.shape[0], 1)  # At least one is not dropped.
+    self.assertEqual(y.shape[0], 1 if words else 0)  # At least one is not dropped.
 
   @parameterized.expand([[1], [2], [4], [5]])
   def testWordOmission(self, count):
Repository: OpenNMT/OpenNMT-tf
Pull request #615 (2020-02-14 12:45:53)

Enable CTranslate2 export of TransformerRelative models


Patch:
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index 0f3f9176..67a3892d 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -95,11 +95,13 @@ class Transformer(SequenceToSequence):
     self._num_units = num_units
     self._num_layers = num_layers
     self._num_heads = num_heads
-    self._is_vanilla_transformer = (
+    self._with_relative_position = maximum_relative_position is not None
+    self._is_ct2_compatible = (
         isinstance(encoder, SelfAttentionEncoder)
         and ffn_activation == tf.nn.relu
-        and position_encoder_class == SinusoidalPositionEncoder
-        and maximum_relative_position is None)
+        and ((self._with_relative_position and position_encoder_class is None)
+             or (not self._with_relative_position
+                 and position_encoder_class == SinusoidalPositionEncoder)))
     super(Transformer, self).__init__(
         source_inputter,
         target_inputter,
@@ -109,10 +111,13 @@ class Transformer(SequenceToSequence):
 
   @property
   def ctranslate2_spec(self):
-    if not self._is_vanilla_transformer:
+    if not self._is_ct2_compatible:
       return None
     import ctranslate2  # pylint: disable=import-outside-toplevel
-    return ctranslate2.specs.TransformerSpec(self._num_layers, self._num_heads)
+    return ctranslate2.specs.TransformerSpec(
+        self._num_layers,
+        self._num_heads,
+        with_relative_position=self._with_relative_position)
 
   def auto_config(self, num_replicas=1):
     config = super(Transformer, self).auto_config(num_replicas=num_replicas)
diff --git a/setup.py b/setup.py
index d453b2cf..32f104e5 100644
--- a/setup.py
+++ b/setup.py
@@ -42,7 +42,7 @@ setup(
     keywords="tensorflow opennmt nmt neural machine translation",
     python_requires='>=3.5',
     install_requires=[
-        "ctranslate2>=1.4,<2;platform_system=='Linux'",
+        "ctranslate2>=1.6,<2;platform_system=='Linux'",
         "pyonmttok>=1.18.1,<2;platform_system=='Linux'",
         "pyyaml>=5.3,<5.4",
         "rouge==0.3.1",


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index eff98c93..538182e6 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -529,6 +529,18 @@ class ModelTest(tf.test.TestCase):
     _check_weight_and_slots(
         lambda model: model.decoder.output_layer.kernel, tgt_mapping, vocab_axis=1)
 
+  @parameterized.expand([
+      [models.TransformerBase()],
+      [models.TransformerBaseRelative()],
+      [models.TransformerBig()],
+      [models.TransformerBigRelative()]
+  ])
+  def testCTranslate2Spec(self, model):
+    try:
+      self.assertIsNotNone(model.ctranslate2_spec)
+    except ImportError:
+      self.skipTest("ctranslate2 module is not available")
+
 
 if __name__ == "__main__":
   tf.test.main()
Repository: OpenNMT/OpenNMT-tf
Pull request #621 (2020-02-25 09:10:37)

Fix decoder initialization with multi source encoder


Patch:
diff --git a/opennmt/decoders/decoder.py b/opennmt/decoders/decoder.py
index 9febf48c..905fa594 100644
--- a/opennmt/decoders/decoder.py
+++ b/opennmt/decoders/decoder.py
@@ -165,13 +165,12 @@ class Decoder(tf.keras.layers.Layer):
     self.memory = memory
     self.memory_sequence_length = memory_sequence_length
     if batch_size is None or dtype is None:
-      if initial_state is None and memory is None:
+      sentinel = tf.nest.flatten(memory)[0]
+      if sentinel is None:
+        sentinel = tf.nest.flatten(initial_state)[0]
+      if sentinel is None:
         raise ValueError("If batch_size or dtype are not set, then either "
                          "memory or initial_state should be set")
-      template = initial_state
-      if template is None:
-        template = memory
-      sentinel = tf.nest.flatten(template)[0]
       if batch_size is None:
         batch_size = tf.shape(sentinel)[0]
       if dtype is None:


Test patch:
diff --git a/opennmt/tests/decoder_test.py b/opennmt/tests/decoder_test.py
index 9b15daf2..7d17242f 100644
--- a/opennmt/tests/decoder_test.py
+++ b/opennmt/tests/decoder_test.py
@@ -17,10 +17,12 @@ def _generate_source_context(batch_size,
   memory = [
       tf.random.uniform([batch_size, time, depth], dtype=dtype)
       for time in memory_time]
+  initial_state = tuple(None for _ in range(num_sources))
   if num_sources == 1:
     memory_sequence_length = memory_sequence_length[0]
     memory = memory[0]
-  return memory, memory_sequence_length
+    initial_state = initial_state[0]
+  return memory, memory_sequence_length, initial_state
 
 
 class DecoderTest(tf.test.TestCase):
@@ -66,13 +68,12 @@ class DecoderTest(tf.test.TestCase):
     vocab_size = 10
     time_dim = 5
     depth = 6
-    memory, memory_sequence_length = _generate_source_context(
+    memory, memory_sequence_length, initial_state = _generate_source_context(
         batch_size,
         depth,
         num_sources=num_sources,
         dtype=dtype)
 
-    initial_state = None
     if initial_state_fn is not None:
       initial_state = initial_state_fn(batch_size, dtype)
     decoder.initialize(vocab_size=vocab_size)
Repository: OpenNMT/OpenNMT-tf
Pull request #624 (2020-02-27 16:02:03)

Log more information on length mismatch when parsing alignments


Patch:
diff --git a/opennmt/data/text.py b/opennmt/data/text.py
index 9ab47e5f..e4b32ef1 100644
--- a/opennmt/data/text.py
+++ b/opennmt/data/text.py
@@ -90,7 +90,22 @@ def alignment_matrix_from_pharaoh(alignment_line,
   sparse_values = tf.ones([tf.shape(sparse_indices)[0]], dtype=dtype)
   source_length = tf.cast(source_length, tf.int64)
   target_length = tf.cast(target_length, tf.int64)
-  alignment_matrix_sparse = tf.sparse.SparseTensor(
-      sparse_indices, sparse_values, [source_length, target_length])
-  alignment_matrix = tf.sparse.to_dense(alignment_matrix_sparse, validate_indices=False)
-  return tf.transpose(alignment_matrix)
+  maximum_ids = tf.reduce_max(sparse_indices, axis=0)
+  assert_source_length = _assert_in_range(maximum_ids[0], source_length, alignment_line, "source")
+  assert_target_length = _assert_in_range(maximum_ids[1], target_length, alignment_line, "target")
+  with tf.control_dependencies([assert_source_length, assert_target_length]):
+    alignment_matrix_sparse = tf.sparse.SparseTensor(
+        sparse_indices, sparse_values, [source_length, target_length])
+    alignment_matrix = tf.sparse.to_dense(alignment_matrix_sparse, validate_indices=False)
+    return tf.transpose(alignment_matrix)
+
+
+def _assert_in_range(maximum_id, length, line, name):
+  return tf.debugging.assert_less(
+      maximum_id,
+      length,
+      message=tf.strings.format(
+          "Length mismatch for alignment line {}: actual %s length is {}, but "
+          "got %s id {} which is out of range. Please check that the alignment "
+          "file is correctly aligned to the training file." % (name, name),
+          [line, length, maximum_id]))


Test patch:
diff --git a/opennmt/tests/text_test.py b/opennmt/tests/text_test.py
index e7cb46ce..ec579880 100644
--- a/opennmt/tests/text_test.py
+++ b/opennmt/tests/text_test.py
@@ -63,6 +63,17 @@ class TextTest(tf.test.TestCase):
     self._testPharaohAlignments(
         "0-0 1-2 1-1 2-4", [3, 5], [[1, 0, 0], [0, 1, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]])
 
+  @parameterized.expand([[True], [False]])
+  def testInvalidPharaohAlignments(self, run_as_function):
+    func = text.alignment_matrix_from_pharaoh
+    if run_as_function:
+      func = tf.function(func)
+
+    with self.assertRaisesRegex(tf.errors.InvalidArgumentError, "source"):
+      func(tf.constant("0-0 1-1 2-3 3-2"), 2, 4)
+    with self.assertRaisesRegex(tf.errors.InvalidArgumentError, "target"):
+      func(tf.constant("0-0 1-2 1-1 2-4"), 3, 4)
+
 
 if __name__ == "__main__":
   tf.test.main()
Repository: OpenNMT/OpenNMT-tf
Pull request #625 (2020-02-28 14:01:55)

Enable CTranslate2 export for Transformers with different enc/dec layers


Patch:
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index 58970613..b4c3cf4e 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -104,7 +104,6 @@ class Transformer(SequenceToSequence):
     self._with_relative_position = maximum_relative_position is not None
     self._is_ct2_compatible = (
         isinstance(encoder, SelfAttentionEncoder)
-        and num_encoder_layers == num_decoder_layers
         and ffn_activation == tf.nn.relu
         and ((self._with_relative_position and position_encoder_class is None)
              or (not self._with_relative_position
@@ -122,7 +121,7 @@ class Transformer(SequenceToSequence):
       return None
     import ctranslate2  # pylint: disable=import-outside-toplevel
     return ctranslate2.specs.TransformerSpec(
-        self._num_encoder_layers,
+        (self._num_encoder_layers, self._num_decoder_layers),
         self._num_heads,
         with_relative_position=self._with_relative_position)
 
diff --git a/setup.py b/setup.py
index 32f104e5..aa07ad28 100644
--- a/setup.py
+++ b/setup.py
@@ -42,7 +42,7 @@ setup(
     keywords="tensorflow opennmt nmt neural machine translation",
     python_requires='>=3.5',
     install_requires=[
-        "ctranslate2>=1.6,<2;platform_system=='Linux'",
+        "ctranslate2>=1.7,<2;platform_system=='Linux'",
         "pyonmttok>=1.18.1,<2;platform_system=='Linux'",
         "pyyaml>=5.3,<5.4",
         "rouge==0.3.1",


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 3932d225..109fa21e 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -533,7 +533,14 @@ class ModelTest(tf.test.TestCase):
       [models.TransformerBase()],
       [models.TransformerBaseRelative()],
       [models.TransformerBig()],
-      [models.TransformerBigRelative()]
+      [models.TransformerBigRelative()],
+      [models.Transformer(
+          inputters.WordEmbedder(32),
+          inputters.WordEmbedder(32),
+          num_layers=(6, 3),
+          num_units=32,
+          num_heads=8,
+          ffn_inner_dim=64)],
   ])
   def testCTranslate2Spec(self, model):
     try:
Repository: OpenNMT/OpenNMT-tf
Pull request #626 (2020-03-02 14:51:57)

Prepare for primary attribute of DistributedValues going private
cf. https://github.com/tensorflow/tensorflow/commit/19ac5f4f6c44ce98654f26c24bb8cd3971c821ab


Patch:
diff --git a/opennmt/training.py b/opennmt/training.py
index b550aac1..16418131 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -402,7 +402,7 @@ class MovingAverage(object):
   @tf.function
   def update(self):
     """Updates the moving average of the variables."""
-    self._ema.apply(var_list=list(map(_get_primary_variable, self._variables)))
+    self._ema.apply(var_list=list(map(misc.get_primary_variable, self._variables)))
 
   @contextlib.contextmanager
   def shadow_variables(self):
@@ -416,14 +416,7 @@ class MovingAverage(object):
     previous_values = []
     for variable in self._variables:
       previous_values.append(variable.value())
-      variable.assign(self._ema.average(_get_primary_variable(variable)))
+      variable.assign(self._ema.average(misc.get_primary_variable(variable)))
     yield
     for previous_value, variable in zip(previous_values, self._variables):
       variable.assign(previous_value)
-
-
-def _get_primary_variable(variable):
-  """Returns the primary variable of a MirroredVariable if possible, else the
-  variable itself.
-  """
-  return getattr(variable, "primary", variable)
diff --git a/opennmt/utils/misc.py b/opennmt/utils/misc.py
index d2ce01c4..405abcf2 100644
--- a/opennmt/utils/misc.py
+++ b/opennmt/utils/misc.py
@@ -64,9 +64,18 @@ def get_variable_name(variable, root, model_key="model"):
   """Gets the variable name in the object-based representation."""
   variables_to_names, _ = get_variables_name_mapping(root, root_key=model_key)
   # In case of a MirroredVariable, look up the primary variable
-  variable = getattr(variable, "primary", variable)
+  variable = get_primary_variable(variable)
   return variables_to_names.get(variable.experimental_ref())
 
+def get_primary_variable(variable):
+  """If :obj:`variable` is distributed, returns the primary component."""
+  # TODO: use a public API to get the primary variable.
+  for attribute_name in ("primary", "_primary"):
+    attribute = getattr(variable, attribute_name, None)
+    if attribute is not None:
+      return attribute
+  return variable
+
 def print_as_bytes(text, stream=None):
   """Prints a string as bytes to non rely on :obj:`stream` default encoding.
 


Test patch:
diff --git a/opennmt/tests/training_test.py b/opennmt/tests/training_test.py
index c906b24a..aca5dfab 100644
--- a/opennmt/tests/training_test.py
+++ b/opennmt/tests/training_test.py
@@ -33,7 +33,6 @@ class TrainingTest(tf.test.TestCase):
     moving_average.update()
     self.assertAllEqual(self.evaluate(variables), [3.0, 4.0])
     with moving_average.shadow_variables():
-      print(variables)
       self.assertAllClose(self.evaluate(variables), [2.8, 3.8])
     self.assertAllEqual(self.evaluate(variables), [3.0, 4.0])
Repository: OpenNMT/OpenNMT-tf
Pull request #659 (2020-05-06 08:01:40)

Fix type error when tokenizing an empty tensor


Patch:
diff --git a/opennmt/tokenizers/tokenizer.py b/opennmt/tokenizers/tokenizer.py
index 23b9cb88..f7b817c6 100644
--- a/opennmt/tokenizers/tokenizer.py
+++ b/opennmt/tokenizers/tokenizer.py
@@ -159,7 +159,7 @@ class Tokenizer(abc.ABC):
     def _python_wrapper(string_t):
       string = tf.compat.as_text(string_t.numpy())
       tokens = self._tokenize_string(string)
-      return tf.constant(tokens)
+      return tf.constant(tokens, dtype=tf.string)
     tokens = tf.py_function(_python_wrapper, [text], tf.string)
     tokens.set_shape([None])
     return tokens


Test patch:
diff --git a/opennmt/tests/tokenizer_test.py b/opennmt/tests/tokenizer_test.py
index 869985cb..bf3d627a 100644
--- a/opennmt/tests/tokenizer_test.py
+++ b/opennmt/tests/tokenizer_test.py
@@ -103,6 +103,15 @@ class TokenizerTest(tf.test.TestCase):
         [["Hello", "world", "￭!"], ["Test"], ["My", "name"]],
         ["Hello world!", "Test", "My name"])
 
+  def testOpenNMTTokenizerEmptyTensor(self):
+    tokenizer = tokenizers.OpenNMTTokenizer()
+    tokens = tokenizer.tokenize(tf.constant(""))
+    self.assertIs(tokens.dtype, tf.string)
+    self.assertListEqual(tokens.shape.as_list(), [0])
+    text = tokenizer.detokenize(tokens)
+    self.assertIs(text.dtype, tf.string)
+    self.assertListEqual(text.shape.as_list(), [])
+
   def testOpenNMTTokenizerArguments(self):
     tokenizer = tokenizers.OpenNMTTokenizer(
         mode="aggressive", spacer_annotate=True, spacer_new=True)
Repository: OpenNMT/OpenNMT-tf
Pull request #729 (2020-10-19 10:38:50)

Make error more explicit when the input layer is not initialized


Patch:
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index bf7bb975..71d35d0f 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -218,6 +218,11 @@ class TextInputter(Inputter):
     self.noiser = None
     self.in_place_noise = False
     self.noise_probability = None
+    self.vocabulary_file = None
+    self.vocabulary_size = None
+    self.tokens_to_ids = None
+    self.ids_to_tokens = None
+    self.tokenizer = None
 
   def initialize(self, data_config, asset_prefix=""):
     self.vocabulary_file = _get_field(
@@ -250,6 +255,7 @@ class TextInputter(Inputter):
     self.noise_probability = probability
 
   def export_assets(self, asset_dir, asset_prefix=""):
+    self._assert_is_initialized()
     return self.tokenizer.export_assets(asset_dir, asset_prefix=asset_prefix)
 
   def make_dataset(self, data_file, training=None):
@@ -262,6 +268,7 @@ class TextInputter(Inputter):
 
   def make_features(self, element=None, features=None, training=None):
     """Tokenizes raw text."""
+    self._assert_is_initialized()
     if features is None:
       features = {}
     if "tokens" in features:
@@ -292,6 +299,7 @@ class TextInputter(Inputter):
     return features
 
   def input_signature(self):
+    self._assert_is_initialized()
     if self.tokenizer.in_graph:
       return {
           "text": tf.TensorSpec([None], tf.string)
@@ -302,6 +310,11 @@ class TextInputter(Inputter):
           "length": tf.TensorSpec([None], tf.int32)
       }
 
+  def _assert_is_initialized(self):
+    if self.tokenizer is None:
+      raise RuntimeError("The input layer is not initialized. You should initialize "
+                         "the model by calling model.initialize(data_config).")
+
 
 class WordEmbedder(TextInputter):
   """Simple word embedder."""


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index dd20abe2..fe87f222 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -303,6 +303,13 @@ class InputterTest(tf.test.TestCase):
     self.assertAllEqual([1, 1], transformed[0][0])
     self.assertAllEqual([2, 2], transformed[0][1])
 
+  def testWordEmbedderMissingInitialization(self):
+    embedder = text_inputter.WordEmbedder()
+    with self.assertRaisesRegex(RuntimeError, "initialize"):
+      embedder.input_signature()
+    with self.assertRaisesRegex(RuntimeError, "initialize"):
+      embedder.make_features("Hello world !")
+
   def testCharConvEmbedder(self):
     vocab_file = self._makeTextFile("vocab.txt", ["h", "e", "l", "w", "o"])
     data_file = self._makeTextFile("data.txt", ["hello world !"])
Repository: OpenNMT/OpenNMT-tf
Pull request #758 (2020-12-15 07:57:17)

Ensure training summary is serializable to JSON


Patch:
diff --git a/opennmt/training.py b/opennmt/training.py
index 3c6b2a59..dd238948 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -595,9 +595,7 @@ class TrainingStats:
     learning_rate = self._optimizer.learning_rate
     if isinstance(learning_rate, tf.optimizers.schedules.LearningRateSchedule):
       learning_rate = learning_rate(self._last_step)
-    elif isinstance(learning_rate, tf.Variable):
-      learning_rate = learning_rate.value()
-    return learning_rate
+    return float(learning_rate)
 
   def _update_words_counter(self, name, features):
     """Accumulates the number of source and target tokens to report throughput."""


Test patch:
diff --git a/opennmt/tests/training_test.py b/opennmt/tests/training_test.py
index f625f1ab..6189df44 100644
--- a/opennmt/tests/training_test.py
+++ b/opennmt/tests/training_test.py
@@ -1,4 +1,5 @@
 import os
+import json
 
 import tensorflow as tf
 
@@ -84,10 +85,15 @@ class TrainingTest(tf.test.TestCase):
       stats.update_on_example(source_features, target_features)
       stats.update_on_step(step, loss)
 
+    def _is_json_serializable(summary):
+      json.dumps(summary)
+      return True
+
     _step(24, 23, 5, 9.8)
     _step(10, 8, 10, 9.6)
 
     summary = stats.get_last_summary()
+    self.assertTrue(_is_json_serializable(summary))
     self.assertEqual(summary["learning_rate"], 1.0)
     self.assertEqual(summary["step"], 10)
     self.assertEqual(summary["loss"], 9.6)
@@ -113,6 +119,7 @@ class TrainingTest(tf.test.TestCase):
     self.assertEqual(summary["words_per_sec"]["target"], 0)
 
     summary = stats.get_global_summary()
+    self.assertTrue(_is_json_serializable(summary))
     self.assertEqual(summary["last_learning_rate"], 1.0)
     self.assertEqual(summary["last_step"], 15)
     self.assertEqual(summary["last_loss"], 9.4)