import tensorflow as tf

      out = spinn._bundle(lstm_iter)

      self.assertEqual(2, len(out))
      self.assertEqual(tf.float32, out[0].dtype)
      self.assertEqual(tf.float32, out[1].dtype)
      self.assertAllEqual(np.array([[0, 2, 0, -2, 0, 4, 0, -4]]).T,
                          out[0].numpy())
      self.assertAllEqual(np.array([[1, 3, -1, -3, 2, 6, -2, -6]]).T,
                          out[1].numpy())

  def testUnbunbdle(self):
    with tf.device(self._test_device):
      state = [np.array([[0, 1, 2], [3, 4, 5]], dtype=np.float32),
               np.array([[0, -1, -2], [-3, -4, -5]], dtype=np.float32)]
      out = spinn._unbundle(state)

      self.assertEqual(2, len(out))
      self.assertEqual(tf.float32, out[0].dtype)
      self.assertEqual(tf.float32, out[1].dtype)
      self.assertAllEqual(np.array([[0, 1, 2, 0, -1, -2]]),
                          out[0].numpy())
      self.assertAllEqual(np.array([[3, 4, 5, -3, -4, -5]]),
