value_extraction_by_index

x = np.array([[0.49593696, 0.504063  ],
              [0.4912244 , 0.50877565],
              [0.48871803, 0.51128197],
              [0.48469874, 0.5153013 ],
              [0.4801116 , 0.5198884 ]])
a = np.array([0,0,1,1,0]).astype(np.int32)

# numpy array
Z = x[range(5),a] 

#==========================================
# tensorflow
X = tf.convert_to_tensor(x)
A = tf.convert_to_tensor(a)

w = tf.stack([tf.range(tf.shape(X)[0]),A],axis=-1)  # tf.shape(X)[0] <--- batch_size. 여기서는 5
# w: array([[0, 0],[1, 0],[2, 1],[3, 1],[4, 0]])
z1 = tf.gather_nd(X,w)  # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법
ww = tf.range(0, tf.shape(X)[0]) * tf.shape(X)[1] + A
# [0,1,2,3,4]  --> [0,2,4,6,8] --> [0,2,5,7,8]
z2 = tf.gather(tf.reshape(X, [-1]), ww)  # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법2
one_hot_A = tf.one_hot(A,tf.shape(X)[1])
z3 = tf.reduce_sum(X*one_hot_A,axis=-1) # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법3
z4 = tf.boolean_mask(X,one_hot_A)

z5 = tf.exp(-tf.nn.softmax_cross_entropy_with_logits_v2(logits=tf.log(X),labels=one_hot_A, dim=-1))

sess = tf.Session()

sess.run([z1,z2,z3,z4,z5])



# pytorch code
X = torch.from_numpy(x)
A = torch.from_numpy(a).long()

z = torch.gather(X,1,A.view(-1,1)).squeeze(1)  # A: 5x1로 만들어 넣어주면, 5x1 크기가 return 되어 온다. 
z = torch.gather(X,1,A.unsqueeze(1)).squeeze(1) # 이렇게 해도 된다.
z = X[range(5),a]

placeholder version

x = np.array([[0.49593696, 0.504063  ],
              [0.4912244 , 0.50877565],
              [0.48871803, 0.51128197],
              [0.48469874, 0.5153013 ],
              [0.4801116 , 0.5198884 ]])
a = np.array([0,0,1,1,0]).astype(np.int32)

# numpy array
Z = x[range(5),a] 

#==========================================
# tensorflow
X = tf.placeholder(tf.float32,[None,2]) 
A = tf.placeholder(tf.int32,[None]) 


w = tf.stack([tf.range(tf.shape(X)[0]),A],axis=-1)  # tf.shape(X)[0] <--- batch_size. 여기서는 5
# w: array([[0, 0],[1, 0],[2, 1],[3, 1],[4, 0]])
z1 = tf.gather_nd(X,w)  # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법
ww = tf.range(0, tf.shape(X)[0]) * tf.shape(X)[1] + A
# [0,1,2,3,4]  --> [0,2,4,6,8] --> [0,2,5,7,8]
z2 = tf.gather(tf.reshape(X, [-1]), ww)  # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법2
one_hot_A = tf.one_hot(A,tf.shape(X)[1])
z3 = tf.reduce_sum(X*one_hot_A,axis=-1) # ---> array([0.49593696, 0.4912244 , 0.51128197, 0.5153013 , 0.4801116 ]

# 다른 방법3
z4 = tf.boolean_mask(X,one_hot_A)

z5 = tf.exp(-tf.nn.softmax_cross_entropy_with_logits_v2(logits=tf.log(X),labels=one_hot_A, dim=-1))

sess = tf.Session()

sess.run([z1,z2,z3,z4,z5], feed_dict={X: x[:3], A: a[:3]})