import tensorflow as tf

  def testSPINN(self):
    with tf.device(self._test_device):
      embedding_dims = 10
      d_tracker = 8
      sequence_length = 15
      num_transitions = 27

      config_tuple = collections.namedtuple(
          "Config", ["d_hidden", "d_proj", "d_tracker", "predict"])
      config = config_tuple(
          embedding_dims, embedding_dims * 2, d_tracker, False)
      s = spinn.SPINN(config)

      # Create some fake data.
      buffers = tf.random_normal((sequence_length, 1, config.d_proj))
      transitions = tf.constant(
          [[3], [3], [2], [3], [3], [3], [2], [2], [2], [3], [3], [3],
           [2], [3], [3], [2], [2], [3], [3], [3], [2], [2], [2], [2],
           [3], [2], [2]], dtype=tf.int64)
      self.assertEqual(tf.int64, transitions.dtype)
      self.assertEqual((num_transitions, 1), transitions.shape)

      out = s(buffers, transitions, training=True)
      self.assertEqual(tf.float32, out.dtype)
      self.assertEqual((1, embedding_dims), out.shape)

  def testSNLIClassifierAndTrainer(self):
    with tf.device(self._test_device):
