Cherry-pick index sample op in contrib (#23522)

* test=develop, add index sample op in contrib

paddle/fluid/operators/index_sample_op.cc

+/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/index_sample_op.h"
+#include <memory>
+#include <vector>
+#include "paddle/fluid/framework/no_need_buffer_vars_inference.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace operators {
+class IndexSampleOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X", "Input(Tensor), dtype support int32/int64/float/double");
+    AddInput("Index", "Index(Tensor), dtype support int32/int64");
+    AddOutput("Out", "Return the element of input at index");
+
+    AddComment(R"DOC(
+    IndexSample OP returns the element of the specified location of X, 
+    and the location is specified by Index. 
+
+    X tensor and Index tensor's shape must be 2-D, 
+    dimension at 0 which usually is batch size must be equal.
+
+    The returned tensor has the same shape and dimensions as the Index tensor.
+    )DOC");
+  }
+};
+
+class IndexSampleOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
+                      platform::errors::InvalidArgument(
+                          "Inputs(Input) of FindByIndex should not be null."));
+    PADDLE_ENFORCE_EQ(ctx->HasInput("Index"), true,
+                      platform::errors::InvalidArgument(
+                          "Inputs(Index) of FindByIndex should not be null."));
+
+    auto input_dims = ctx->GetInputDim("X");
+    PADDLE_ENFORCE_EQ(
+        input_dims.size(), 2,
+        platform::errors::InvalidArgument(
+            "Inputs(X) shape of IndexSample op should be 2-D, but "
+            "got X's shape = [%s], please check X shape.",
+            input_dims));
+
+    auto index_dims = ctx->GetInputDim("Index");
+    PADDLE_ENFORCE_EQ(
+        input_dims.size(), 2,
+        platform::errors::InvalidArgument(
+            "Inputs(Index) shape of IndexSample op should be 2-D, but "
+            "got Index's shape [%s] , please check index shape.",
+            input_dims));
+    if (ctx->IsRuntime()) {
+      PADDLE_ENFORCE_EQ(input_dims[0], index_dims[0],
+                        platform::errors::InvalidArgument(
+                            "Inputs(X)'s value of dimension 0 must same with "
+                            "Inputs(Index)'s value of dimension 0, but "
+                            "got %d of Inputs(X), and got %d of Inputs(Index), "
+                            "please check Inputs shape.",
+                            input_dims[0], index_dims[0]));
+    }
+    ctx->SetOutputDim("Out", index_dims);
+    auto type = ctx->GetInputsVarType("Index")[0];
+    if (type == framework::proto::VarType::LOD_TENSOR) {
+      ctx->ShareLoD("Index", /*->*/ "Out");
+    }
+  }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
+    return framework::OpKernelType(data_type, ctx.device_context());
+  }
+};
+
+class IndexSampleGradOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    PADDLE_ENFORCE_EQ(
+        ctx->HasInput("Index"), true,
+        platform::errors::InvalidArgument("Input(Index) should be not null."));
+    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
+                      platform::errors::InvalidArgument(
+                          "Input(Out@GRAD) should be not null."));
+    PADDLE_ENFORCE_EQ(ctx->HasOutput(framework::GradVarName("X")), true,
+                      platform::errors::InvalidArgument(
+                          "Output(X@GRAD) should be not null."));
+
+    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
+  }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    auto data_type = OperatorWithKernel::IndicateVarDataType(
+        ctx, framework::GradVarName("Out"));
+    return framework::OpKernelType(data_type, ctx.device_context());
+  }
+};
+
+template <typename T>
+class IndexSampleGradMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  std::unique_ptr<T> Apply() const override {
+    std::unique_ptr<T> op(new T());
+    op->SetType("index_sample_grad");
+    op->SetInput("X", this->Input("X"));
+    op->SetInput("Index", this->Input("Index"));
+    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
+    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
+    return op;
+  }
+};
+
+DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(IndexSampleGradNoNeedBufferVarInferer,
+                                      "X");
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OPERATOR(index_sample, ops::IndexSampleOp, ops::IndexSampleOpMaker,
+                  ops::IndexSampleGradMaker<paddle::framework::OpDesc>,
+                  ops::IndexSampleGradMaker<paddle::imperative::OpBase>);
+REGISTER_OPERATOR(index_sample_grad, ops::IndexSampleGradOp,
+                  ops::IndexSampleGradNoNeedBufferVarInferer);
+REGISTER_OP_CPU_KERNEL(
+    index_sample, ops::IndexSampleKernel<paddle::platform::CPUPlace, float>,
+    ops::IndexSampleKernel<paddle::platform::CPUPlace, double>,
+    ops::IndexSampleKernel<paddle::platform::CPUPlace, int>,
+    ops::IndexSampleKernel<paddle::platform::CPUPlace, int64_t>);
+REGISTER_OP_CPU_KERNEL(
+    index_sample_grad,
+    ops::IndexSampleGradKernel<paddle::platform::CPUPlace, float>,
+    ops::IndexSampleGradKernel<paddle::platform::CPUPlace, double>,
+    ops::IndexSampleGradKernel<paddle::platform::CPUPlace, int>,
+    ops::IndexSampleGradKernel<paddle::platform::CPUPlace, int64_t>);

paddle/fluid/operators/index_sample_op.h

+/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <gflags/gflags.h>
+#include <cmath>
+#include <fstream>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+#include "paddle/fluid/framework/op_registry.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+using LoDTensor = framework::LoDTensor;
+using DDim = framework::DDim;
+
+template <typename T, typename IndexT = int>
+void IndexSampleInner(const framework::ExecutionContext &context,
+                      const LoDTensor &input, const LoDTensor &index,
+                      LoDTensor *output) {
+  auto input_dims = input.dims();
+  auto index_dims = index.dims();
+
+  int batch_size = input_dims[0];
+  auto value_length = input_dims[1];
+  auto index_length = index_dims[1];
+  int index_ids_num = index.numel();
+  auto *input_data = input.data<T>();
+  auto *index_data = index.data<IndexT>();
+
+  std::vector<T> res{};
+  for (int i = 0; i < index_ids_num; i++) {
+    int b = floor(i / index_length);
+    PADDLE_ENFORCE_GE(
+        index_data[i], 0,
+        platform::errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            value_length, index_data[i]));
+    PADDLE_ENFORCE_LT(
+        index_data[i], value_length,
+        platform::errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            value_length, index_data[i]));
+
+    int v_i = b * value_length + static_cast<int>(index_data[i]);
+    T v = input_data[v_i];
+    VLOG(4) << "Index Sample: batch = " << b << " index = " << v_i
+            << " value = " << v;
+    res.push_back(v);
+  }
+
+  auto ddim = framework::make_ddim({batch_size, index_length});
+  output->Resize(ddim);
+  T *out_data = output->mutable_data<T>(context.GetPlace());
+
+  memcpy(out_data, &res[0], sizeof(T) * index_ids_num);
+}
+
+template <typename DeviceContext, typename T>
+class IndexSampleKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext &ctx) const override {
+    auto *input_var = ctx.InputVar("X");
+    auto *index_var = ctx.InputVar("Index");
+
+    auto &input_tensor = input_var->Get<LoDTensor>();
+    auto &index_tensor = index_var->Get<LoDTensor>();
+
+    auto *out_var = ctx.OutputVar("Out");
+    auto *out_tensor = out_var->GetMutable<framework::LoDTensor>();
+
+    const auto &index_type = index_tensor.type();
+    bool index_type_match = index_type == framework::proto::VarType::INT32 ||
+                            index_type == framework::proto::VarType::INT64;
+    PADDLE_ENFORCE_EQ(index_type_match, true,
+                      platform::errors::InvalidArgument(
+                          "Input(Index) holds the wrong type, it holds %s, but "
+                          "desires to be %s or %s",
+                          paddle::framework::DataTypeToString(index_type),
+                          paddle::framework::DataTypeToString(
+                              framework::proto::VarType::INT32),
+                          paddle::framework::DataTypeToString(
+                              framework::proto::VarType::INT64)));
+    if (index_type == framework::proto::VarType::INT32) {
+      IndexSampleInner<T, int>(ctx, input_tensor, index_tensor, out_tensor);
+    } else if (index_type == framework::proto::VarType::INT64) {
+      IndexSampleInner<T, int64_t>(ctx, input_tensor, index_tensor, out_tensor);
+    }
+  }
+};
+
+template <typename T, typename IndexT = int>
+void IndexSampleGradInner(const framework::ExecutionContext &context,
+                          const LoDTensor &out_grad, const LoDTensor &index,
+                          LoDTensor *x_grad) {
+  auto index_dims = index.dims();
+  auto x_grad_dims = x_grad->dims();
+
+  int batch_size = x_grad_dims[0];
+  auto value_length = x_grad_dims[1];
+  auto index_length = index_dims[1];
+  int index_ids_num = index.numel();
+
+  T *x_grad_data = x_grad->mutable_data<T>(context.GetPlace());
+  auto *out_grad_data = out_grad.data<T>();
+  auto *index_data = index.data<IndexT>();
+
+  memset(x_grad_data, 0, batch_size * value_length * sizeof(T));
+
+  for (int i = 0; i < index_ids_num; i++) {
+    int b = floor(i / index_length);
+    PADDLE_ENFORCE_GE(
+        index_data[i], 0,
+        platform::errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample_grad) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            value_length, index_data[i]));
+    PADDLE_ENFORCE_LT(
+        index_data[i], value_length,
+        platform::errors::InvalidArgument(
+            "Variable value (index) of OP(index_sample_grad) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            value_length, index_data[i]));
+    int v_i = b * value_length + static_cast<int>(index_data[i]);
+    x_grad_data[v_i] += out_grad_data[i];
+  }
+}
+
+template <typename DeviceContext, typename T>
+class IndexSampleGradKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext &context) const override {
+    auto *index_var = context.InputVar("Index");
+    auto *x_grad_var = context.OutputVar(framework::GradVarName("X"));
+    auto *out_grad_var = context.InputVar(framework::GradVarName("Out"));
+
+    auto &index_tensor = index_var->Get<LoDTensor>();
+    auto &out_grad_tensor = out_grad_var->Get<LoDTensor>();
+    auto *x_grad_tensor = x_grad_var->GetMutable<framework::LoDTensor>();
+
+    const auto &index_type = index_tensor.type();
+    bool index_type_match = index_type == framework::proto::VarType::INT32 ||
+                            index_type == framework::proto::VarType::INT64;
+    PADDLE_ENFORCE_EQ(index_type_match, true,
+                      platform::errors::InvalidArgument(
+                          "Input(Index) holds the wrong type, it holds %s, but "
+                          "desires to be %s or %s",
+                          paddle::framework::DataTypeToString(index_type),
+                          paddle::framework::DataTypeToString(
+                              framework::proto::VarType::INT32),
+                          paddle::framework::DataTypeToString(
+                              framework::proto::VarType::INT64)));
+    if (index_type == framework::proto::VarType::INT32) {
+      IndexSampleGradInner<T, int>(context, out_grad_tensor, index_tensor,
+                                   x_grad_tensor);
+    } else if (index_type == framework::proto::VarType::INT64) {
+      IndexSampleGradInner<T, int64_t>(context, out_grad_tensor, index_tensor,
+                                       x_grad_tensor);
+    }
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle

python/paddle/fluid/contrib/layers/nn.py

@@ -25,7 +25,7 @@ from paddle.fluid.layer_helper import LayerHelper
 from paddle.fluid.layers import utils
 from ... import unique_name
 from paddle.fluid.initializer import Normal, Constant, NumpyArrayInitializer
-from paddle.fluid.data_feeder import check_type, check_dtype, convert_dtype
+from paddle.fluid.data_feeder import check_type_and_dtype, check_type, check_dtype, convert_dtype
 from paddle.fluid.framework import Variable, convert_np_dtype_to_dtype_
 from paddle.fluid.layers import slice, reshape
 
@@ -39,6 +39,7 @@ __all__ = [
     'multiclass_nms2',
     'search_pyramid_hash',
     'shuffle_batch',
+    'index_sample',
     'tdm_child',
     'tdm_sampler',
 ]
@@ -816,6 +817,66 @@ def shuffle_batch(x, seed=None):
     return out
 
 
+def index_sample(x, index):
+    """
+    **IndexSample Layer**
+    IndexSample OP returns the element of the specified location of X, 
+    and the location is specified by Index. 
+
+    .. code-block:: text
+
+    Given:
+                X = [[1, 2, 3, 4, 5],
+                     [6, 7, 8, 9, 10]]
+                Index = [[0, 1, 3],
+                         [0, 2, 4]]
+                Then:
+                Out = [[1, 2, 4],
+                       [6, 8, 10]]
+
+    Args:
+        x (Variable): The source input tensor with 2-D shape. Supported data type is 
+            int32, int64, float32, float64.
+        index (Variable): The index input tensor with 2-D shape, first dimension should be same with X. 
+            Data type is int32 or int64.
+
+    Returns:
+        output (Variable): The output is a tensor with the same shape as index.
+
+    Examples:
+        .. code-block:: python
+            import paddle.fluid as fluid
+            import numpy as np
+            # create x value
+            x_shape = (2, 5)
+            x_type = "float64"
+            x_np = np.random.random(x_shape).astype(x_type)
+            # create index value
+            index_shape = (2, 3)
+            index_type = "int32"
+            index_np = np.random.randint(low=0, 
+                                         high=x_shape[1],
+                                         size=index_shape).astype(index_type)
+            x = fluid.data(name='x', shape=[-1, 5], dtype='float64')
+            index = fluid.data(name='index', shape=[-1, 3], dtype='int32')
+            output = fluid.contrib.layers.index_sample(x=x, index=index)
+    """
+    helper = LayerHelper("index_sample", **locals())
+    check_type_and_dtype(x, 'x', Variable,
+                         ['float32', 'float64', 'int32', 'int64'],
+                         'fluid.contrib.layers.index_sample')
+    check_type_and_dtype(index, 'index', Variable, ['int32', 'int64'],
+                         'fluid.contrib.layers.index_sample')
+    out = helper.create_variable_for_type_inference(dtype=x.dtype)
+
+    helper.append_op(
+        type='index_sample',
+        inputs={'X': x,
+                'Index': index},
+        outputs={'Out': out})
+    return out
+
+
 def tdm_child(x, node_nums, child_nums, param_attr=None, dtype='int32'):
     """
     **Tdm Child**

python/paddle/fluid/tests/unittests/test_index_sample_op.py

+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import print_function
+
+import unittest
+import numpy as np
+from op_test import OpTest
+
+
+class TestIndexSampleOp(OpTest):
+    def setUp(self):
+        self.op_type = "index_sample"
+        self.config()
+        xnp = np.random.random(self.x_shape).astype(self.x_type)
+        indexnp = np.random.randint(
+            low=0, high=self.x_shape[1],
+            size=self.index_shape).astype(self.index_type)
+        self.inputs = {'X': xnp, 'Index': indexnp}
+        index_array = []
+        for i in range(self.index_shape[0]):
+            for j in indexnp[i]:
+                index_array.append(xnp[i, j])
+        out = np.reshape(index_array, self.index_shape)
+        self.outputs = {'Out': out}
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out')
+
+    def config(self):
+        """
+        For multi-dimension input
+        """
+        self.x_shape = (10, 20)
+        self.x_type = "float64"
+        self.index_shape = (10, 10)
+        self.index_type = "int32"
+
+
+class TestCase1(TestIndexSampleOp):
+    def config(self):
+        """
+        For one dimension input
+        """
+        self.x_shape = (100, 1)
+        self.x_type = "float64"
+        self.index_shape = (100, 1)
+        self.index_type = "int32"
+
+
+class TestCase2(TestIndexSampleOp):
+    def config(self):
+        """
+        For int64_t index type
+        """
+        self.x_shape = (10, 100)
+        self.x_type = "float64"
+        self.index_shape = (10, 10)
+        self.index_type = "int64"
+
+
+class TestCase3(TestIndexSampleOp):
+    def config(self):
+        """
+        For int index type
+        """
+        self.x_shape = (10, 100)
+        self.x_type = "float64"
+        self.index_shape = (10, 10)
+        self.index_type = "int32"
+
+
+class TestCase4(TestIndexSampleOp):
+    def config(self):
+        """
+        For int64 index type
+        """
+        self.x_shape = (10, 100)
+        self.x_type = "float64"
+        self.index_shape = (10, 10)
+        self.index_type = "int64"
+
+
+class TestIndexSampleShape(unittest.TestCase):
+    def test_shape(self):
+        import paddle.fluid as fluid
+        import paddle
+
+        # create x value
+        x_shape = (2, 5)
+        x_type = "float64"
+        x_np = np.random.random(x_shape).astype(x_type)
+
+        # create index value
+        index_shape = (2, 3)
+        index_type = "int32"
+        index_np = np.random.randint(
+            low=0, high=x_shape[1], size=index_shape).astype(index_type)
+
+        x = fluid.data(name='x', shape=[-1, 5], dtype='float64')
+        index = fluid.data(name='index', shape=[-1, 3], dtype='int32')
+        output = fluid.contrib.layers.index_sample(x=x, index=index)
+
+        place = fluid.CPUPlace()
+        exe = fluid.Executor(place=place)
+        exe.run(fluid.default_startup_program())
+
+        feed = {'x': x_np, 'index': index_np}
+        res = exe.run(feed=feed, fetch_list=[output])
+
+
+if __name__ == "__main__":
+    unittest.main()