def testOrdinaryValues(self, keepdims):
    """Indexing value functions by action for a minibatch of values."""
    values = [[1.1, 1.2, 1.3],
              [1.4, 1.5, 1.6],
              [2.1, 2.2, 2.3],
              [2.4, 2.5, 2.6],
              [3.1, 3.2, 3.3],
              [3.4, 3.5, 3.6],
              [4.1, 4.2, 4.3],
              [4.4, 4.5, 4.6]]
    action_indices = [0, 2, 1, 0, 2, 1, 0, 2]
    result = indexing_ops.batched_index(
        values, action_indices, keepdims=keepdims)
    expected_result = [1.1, 1.6, 2.2, 2.4, 3.3, 3.5, 4.1, 4.6]
    if keepdims:
      expected_result = np.expand_dims(expected_result, axis=-1)

    with self.test_session() as sess:
      self.assertAllClose(sess.run(result), expected_result)
Indexing value functions by action for a minibatch of values.
trfl/indexing_ops_test.py