Python _gather_topk_beams Beispiele

Beispiel #1

Datei: beam_search_test.py Projekt: zxfzqf/TransformerTranslationModel

    def test_gather_topk_beams(self):
        x = nd.reshape(nd.arange(24, ctx=ctx), (2, 3, 4))
        x_scores = nd.array([[0, 1, 1], [1, 0, 1]], ctx=ctx)

        y = beam_search._gather_topk_beams([x], x_scores, 2, 2)

        self.assertEqual([[[4, 5, 6, 7], [8, 9, 10, 11]],
                          [[12, 13, 14, 15], [20, 21, 22, 23]]],
                         y[0].asnumpy().tolist())

Beispiel #2

Datei: beam_search_test.py Projekt: mwnewlin/afit_mlperf_training

    def test_gather_topk_beams(self):
        x = tf.reshape(tf.range(24), [2, 3, 4])
        x_scores = [[0, 1, 1], [1, 0, 1]]

        y = beam_search._gather_topk_beams(x, x_scores, 2, 2)
        with self.test_session() as sess:
            y = sess.run(y)

        self.assertAllEqual([[[4, 5, 6, 7], [8, 9, 10, 11]],
                             [[12, 13, 14, 15], [20, 21, 22, 23]]], y)

Beispiel #3

Datei: beam_search_test.py Projekt: cybermaster/reference

  def test_gather_topk_beams(self):
    x = tf.reshape(tf.range(24), [2, 3, 4])
    x_scores = [[0, 1, 1], [1, 0, 1]]

    y = beam_search._gather_topk_beams(x, x_scores, 2, 2)
    with self.test_session() as sess:
      y = sess.run(y)

    self.assertAllEqual([[[4, 5, 6, 7],
                          [8, 9, 10, 11]],
                         [[12, 13, 14, 15],
                          [20, 21, 22, 23]]],
                        y)