Example using tf.gather_nd / tf.scatter_nd

gather_scatter.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow as tf

def gather_scatter_example():
  batch_size = 2
  seq_len = 5
  num_classes = 3
  tf.set_random_seed(23)

  logits = tf.random_normal([batch_size, seq_len, seq_len, num_classes])
  gather_indices = tf.constant([[0, 1, 1],
                                [0, 1, 2], 
                                [0, 2, 3],
                                [1, 0, 1]])

  with tf.Session() as sess:

    g = tf.gather_nd(logits, gather_indices)
    results = sess.run([logits, tf.shape(logits), gather_indices, g])

  for r in results:
    print(r)