Add tests for min and max function when returning indices.

python/test/function/test_reduction.py

@@ -73,3 +73,45 @@ def test_large_mean_reduction(seed, inshape, axis, ctx, func_name):
     inputs = [rng.randn(*inshape).astype(np.float32)]
     function_tester(rng, F.mean, np.mean, inputs, ctx=ctx, func_name=func_name,
                     func_args=[axis])
+
+
+@pytest.mark.parametrize("seed", [313])
+@pytest.mark.parametrize("ctx, func_name", list_context('Max'))
+@pytest.mark.parametrize("keepdims", [False, True])
+@pytest.mark.parametrize("inshape, axis", [
+    ((8, 1, 128, 128), (2, 3)),
+    ((8, 128, 1, 128), (1, 3)),
+])
+def test_max_with_index(seed, ctx, func_name, inshape, axis, keepdims):
+    x = np.random.RandomState(seed).randn(*inshape).astype(np.float32)
+    x = nn.Variable.from_numpy_array(x)
+    with nn.context_scope(ctx), nn.auto_forward(True):
+        val, idx = F.max(x, axis, keepdims, with_index=True)
+    assert np.allclose(val.d, np.amax(x.d, axis, keepdims=keepdims))
+    shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
+    assert np.all(idx.d == x.d.reshape(*shape).argmax(-1).reshape(idx.d.shape))
+    with nn.context_scope(ctx), nn.auto_forward(True):
+        idx = F.max(x, axis, keepdims, only_index=True)
+    shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
+    assert np.all(idx.d == x.d.reshape(*shape).argmax(-1).reshape(idx.d.shape))
+
+
+@pytest.mark.parametrize("seed", [313])
+@pytest.mark.parametrize("ctx, func_name", list_context('Min'))
+@pytest.mark.parametrize("keepdims", [False, True])
+@pytest.mark.parametrize("inshape, axis", [
+    ((8, 1, 128, 128), (2, 3)),
+    ((8, 128, 1, 128), (1, 3)),
+])
+def test_min_with_index(seed, ctx, func_name, inshape, axis, keepdims):
+    x = np.random.RandomState(seed).randn(*inshape).astype(np.float32)
+    x = nn.Variable.from_numpy_array(x)
+    with nn.context_scope(ctx), nn.auto_forward(True):
+        val, idx = F.min(x, axis, keepdims, with_index=True)
+    assert np.allclose(val.d, np.amin(x.d, axis, keepdims=keepdims))
+    shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
+    assert np.all(idx.d == x.d.reshape(*shape).argmin(-1).reshape(idx.d.shape))
+    with nn.context_scope(ctx), nn.auto_forward(True):
+        idx = F.min(x, axis, keepdims, only_index=True)
+    shape = [a for i, a in enumerate(x.d.shape) if i not in axis] + [-1]
+    assert np.all(idx.d == x.d.reshape(*shape).argmin(-1).reshape(idx.d.shape))