Added batch_gather to backend.

keras/backend/tensorflow_backend.py

@@ -996,6 +996,37 @@ def gather(reference, indices):
     return tf.gather(reference, indices)
 
 
+def batch_gather(reference, indices):
+    """Batchwise gathering of row indices.
+
+    The numpy equivalent is `reference[np.arange(batch_size), indices]`, where
+    `batch_size` is the first dimension of the reference tensor.
+
+    # Arguments
+        reference: A tensor with ndim >= 2 of shape.
+          (batch_size, dim1, dim2, ..., dimN)
+        indices: A 1d integer tensor of shape (batch_size) satisfying
+          0 <= i < dim2 for each element i.
+
+    # Returns
+        The selected tensor with shape (batch_size, dim2, ..., dimN).
+
+    # Examples
+        1. If reference is `[[3, 5, 7], [11, 13, 17]]` and indices is `[2, 1]`
+        then the result is `[7, 13]`.
+
+        2. If reference is
+        ```
+          [[[2, 3], [4, 5], [6, 7]],
+           [[10, 11], [12, 13], [16, 17]]]
+        ```
+        and indices is `[2, 1]` then the result is `[[6, 7], [12, 13]]`.
+    """
+    batch_size = shape(reference)[0]
+    indices = tf.stack([tf.range(batch_size), indices], axis=1)
+    return tf.gather_nd(reference, indices)
+
+
 # ELEMENT-WISE OPERATIONS
 
 def _normalize_axis(axis, ndim):

keras/backend/theano_backend.py

@@ -420,6 +420,35 @@ def gather(reference, indices):
     return y
 
 
+def batch_gather(reference, indices):
+    """Batchwise gathering of row indices.
+
+    The numpy equivalent is `reference[np.arange(batch_size), indices]`, where
+    `batch_size` is the first dimension of the reference tensor.
+
+    # Arguments
+        reference: A tensor with ndim >= 2 of shape.
+          (batch_size, dim1, dim2, ..., dimN)
+        indices: A 1d integer tensor of shape (batch_size) satisfying
+          0 <= i < dim2 for each element i.
+
+    # Returns
+        The selected tensor with shape (batch_size, dim2, ..., dimN).
+
+    # Examples
+        1. If reference is `[[3, 5, 7], [11, 13, 17]]` and indices is `[2, 1]`
+        then the result is `[7, 13]`.
+        2. If reference is
+        ```
+          [[[2, 3], [4, 5], [6, 7]],
+           [[10, 11], [12, 13], [16, 17]]]
+        ```
+        and indices is `[2, 1]` then the result is `[[6, 7], [12, 13]]`.
+    """
+    batch_size = shape(reference)[0]
+    return reference[T.arange(batch_size), indices]
+
+
 # ELEMENT-WISE OPERATIONS
 
 

tests/keras/backend/backend_test.py

@@ -231,6 +231,20 @@ def test_gather(self):
             y = K.gather(x, indices)
             assert y._keras_shape == (5, 6, 3, 4)
 
+    @pytest.mark.parametrize('x_np, indices_np', [
+        (np.array([[3, 5, 7], [11, 13, 17]]), np.array([2, 1])),
+        (np.array([[[2, 3], [4, 5], [6, 7]],
+                   [[10, 11], [12, 13], [16, 17]]]), np.array([2, 1])),
+    ])
+    @pytest.mark.parametrize('K', [KTH, KTF], ids=["KTH", "KTF"])
+    def test_batch_gather(self, x_np, indices_np, K):
+        x = K.variable(x_np)
+        indices = K.variable(indices_np, dtype='int32')
+        batch_size = x_np.shape[0]
+        assert_allclose(K.eval(K.batch_gather(x, indices)),
+                        x_np[np.arange(batch_size), indices_np],
+                        rtol=1e-5)
+
     def test_value_manipulation(self):
         val = np.random.random((4, 2))
         xth = KTH.variable(val)