Repository: chainer/chainer

Issue:
select_item method
I need a `select_item` method defined as

```
select_item(x, t) == t.choose(x.T)
```



Discussion:

Resolving patch:
diff --git a/chainer/functions/__init__.py b/chainer/functions/__init__.py
--- a/chainer/functions/__init__.py
+++ b/chainer/functions/__init__.py
@@ -12,6 +12,7 @@
 from chainer.functions.array import concat
 from chainer.functions.array import copy
 from chainer.functions.array import reshape
+from chainer.functions.array import select_item
 from chainer.functions.array import split_axis
 from chainer.functions.array import swapaxes
 from chainer.functions.array import transpose
@@ -71,6 +72,8 @@
 reshape = reshape.reshape
 SplitAxis = split_axis.SplitAxis
 split_axis = split_axis.split_axis
+SelectItem = select_item.SelectItem
+select_item = select_item.select_item
 Swapaxes = swapaxes.Swapaxes
 swapaxes = swapaxes.swapaxes
 Transpose = transpose.Transpose
diff --git a/chainer/functions/array/select_item.py b/chainer/functions/array/select_item.py
new file mode 100644
--- /dev/null
+++ b/chainer/functions/array/select_item.py
@@ -0,0 +1,73 @@
+import numpy
+import six
+
+from chainer import cuda
+from chainer import function
+from chainer.utils import type_check
+
+
+class SelectItem(function.Function):
+
+    """Select elements stored in given indicies."""
+
+    def check_type_forward(self, in_types):
+        type_check.expect(in_types.size() == 2)
+
+        x_type, t_type = in_types
+        type_check.expect(
+            x_type.dtype.kind == 'f',
+            t_type.dtype.kind == 'i',
+            x_type.ndim == 2,
+            t_type.ndim == 1,
+            x_type.shape[0] == t_type.shape[0],
+        )
+
+    def forward_cpu(self, inputs):
+        x, t = inputs
+        return t.choose(x.T),
+
+    def forward_gpu(self, inputs):
+        x, t = inputs
+        y = cuda.elementwise(
+            'S t, raw T x, int32 n_channel',
+            'T y',
+            'y = x[i * n_channel + t]',
+            'getitem_fwd'
+        )(t, x, x.shape[1])
+        return y,
+
+    def backward_cpu(self, inputs, grad_outputs):
+        x, t = inputs
+        gloss = grad_outputs[0]
+        gx = numpy.zeros_like(x)
+        gx[six.moves.range(t.size), t] = gloss
+        return gx, None
+
+    def backward_gpu(self, inputs, grad_outputs):
+        x, t = inputs
+        gloss = grad_outputs[0]
+        gx = cuda.cupy.zeros_like(x)
+        gx = cuda.elementwise(
+            'S t, T gloss, int32 n_channel',
+            'raw T gx',
+            'gx[i * n_channel + t] = gloss',
+            'getitem_bwd'
+        )(t, gloss, x.shape[1], gx)
+        return gx, None
+
+
+def select_item(x, t):
+    """Select elements stored in given indicies.
+
+    This function returns ```t.choose(x.T)```, that means
+    ```y[i] == x[i, t[i]]``` for all ```i```.
+
+    Args:
+        x (Variable): Variable storing arrays.
+        t (Variable): Variable storing index numbers.
+
+    Returns:
+        ~chainer.Variable: Variable that holds ```t```-th element of ```x```.
+
+    """
+    return SelectItem()(x, t)


Test patch:
diff --git a/tests/chainer_tests/functions_tests/array_tests/test_select_item.py b/tests/chainer_tests/functions_tests/array_tests/test_select_item.py
new file mode 100644
--- /dev/null
+++ b/tests/chainer_tests/functions_tests/array_tests/test_select_item.py
@@ -0,0 +1,70 @@
+import unittest
+
+import numpy
+
+import chainer
+from chainer import cuda
+from chainer import functions
+from chainer import gradient_check
+from chainer import testing
+from chainer.testing import attr
+
+
+class TestSelectItem(unittest.TestCase):
+    in_shape = (10, 5)
+    out_shape = (10,)
+
+    def setUp(self):
+        self.x_data = numpy.random.uniform(
+            -1, 1, self.in_shape).astype(numpy.float32)
+        self.t_data = numpy.random.randint(
+            0, 2, self.out_shape).astype(numpy.int32)
+        self.gy_data = numpy.random.uniform(
+            -1, 1, self.out_shape).astype(numpy.float32)
+
+    def check_forward(self, x_data, t_data):
+        x = chainer.Variable(x_data)
+        t = chainer.Variable(t_data)
+        y = functions.select_item(x, t)
+        y_exp = cuda.to_cpu(x_data)[range(t_data.size), cuda.to_cpu(t_data)]
+
+        numpy.testing.assert_equal(cuda.to_cpu(y.data), y_exp)
+
+    def test_forward_cpu(self):
+        self.check_forward(self.x_data, self.t_data)
+
+    @attr.gpu
+    def test_forward_gpu(self):
+        self.check_forward(cuda.to_gpu(self.x_data),
+                           cuda.to_gpu(self.t_data))
+
+    def check_backward(self, x_data, t_data, gy_data):
+        x = chainer.Variable(x_data)
+        t = chainer.Variable(t_data)
+        y = functions.select_item(x, t)
+        y.grad = gy_data
+        y.backward()
+        self.assertEqual(None, t.grad)
+
+        func = y.creator
+        f = lambda: func.forward((x.data, t.data))
+        gx, = gradient_check.numerical_grad(f, (x.data,), (gy_data,), eps=0.01)
+
+        gradient_check.assert_allclose(gx, x.grad)
+
+    def test_backward_cpu(self):
+        self.check_backward(self.x_data, self.t_data, self.gy_data)
+
+    @attr.gpu
+    def test_backward_gpu(self):
+        self.check_backward(cuda.to_gpu(self.x_data),
+                            cuda.to_gpu(self.t_data),
+                            cuda.to_gpu(self.gy_data))
+
+
+class TestSelectItemZeroSize(unittest.TestCase):
+    in_shape = (0, 5)
+    out_shape = (0,)
+
+
+testing.run_module(__name__, __file__)