| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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. |
| */ |
| |
| /*! |
| * \file indexing_op.h |
| * \brief Function definition of indexing operator |
| * \author Bing Xu, Siyi Li, Chi Zhang, Haibin Lin |
| */ |
| #ifndef MXNET_OPERATOR_TENSOR_INDEXING_OP_H_ |
| #define MXNET_OPERATOR_TENSOR_INDEXING_OP_H_ |
| |
| #include <dmlc/logging.h> |
| #include <dmlc/parameter.h> |
| #include <mxnet/operator.h> |
| #include <mxnet/operator_util.h> |
| #include <map> |
| #include <vector> |
| #include <string> |
| #include <utility> |
| #include <algorithm> |
| #include <type_traits> |
| #include "../operator_common.h" |
| #include "../mshadow_op.h" |
| #include "../elemwise_op_common.h" |
| #include "./util/tensor_util-inl.h" |
| #include "../mxnet_op.h" |
| #include "./sort_op.h" |
| #include "./init_op.h" |
| #include "../../engine/openmp.h" |
| #include "../../common/utils.h" |
| #ifdef __CUDACC__ |
| #include "./indexing_op-inl.cuh" |
| #endif |
| |
| namespace mxnet { |
| namespace op { |
| |
| namespace embedding { |
| enum EmbeddingOpInputs { kData, kWeight }; |
| enum EmbeddingOpOutputs { kOut }; |
| enum EmbeddingOpResource { kTempSpace }; |
| } // namespace embedding |
| |
| namespace quantized_embedding { |
| enum QuantizedEmbeddingOpInputs { kData, kWeight, kWeightMin, kWeightMax }; |
| enum QuantizedEmbeddingOpOutputs { kOut, kOutMin, kOutMax }; |
| enum QuantizedEmbeddingOpResource { kTempSpace }; |
| } // namespace quantized_embedding |
| |
| struct EmbeddingParam : public dmlc::Parameter<EmbeddingParam> { |
| index_t input_dim; |
| index_t output_dim; |
| int dtype; |
| bool sparse_grad; |
| DMLC_DECLARE_PARAMETER(EmbeddingParam) { |
| DMLC_DECLARE_FIELD(input_dim).set_lower_bound(1).describe( |
| "Vocabulary size of the input indices."); |
| DMLC_DECLARE_FIELD(output_dim) |
| .set_lower_bound(1) |
| .describe("Dimension of the embedding vectors."); |
| DMLC_DECLARE_FIELD(dtype).set_default(mshadow::kFloat32) |
| MXNET_ADD_ALL_TYPES.describe("Data type of weight."); |
| DMLC_DECLARE_FIELD(sparse_grad) |
| .set_default(false) |
| .describe( |
| "Compute row sparse gradient in the backward calculation. If set to True, " |
| "the grad's storage type is row_sparse."); |
| } |
| void SetAttrDict(std::unordered_map<std::string, std::string>* dict) { |
| std::ostringstream input_dim_s, output_dim_s, dtype_s, sparse_grad_s; |
| input_dim_s << input_dim; |
| output_dim_s << output_dim; |
| dtype_s << dtype; |
| sparse_grad_s << sparse_grad; |
| (*dict)["input_dim"] = input_dim_s.str(); |
| (*dict)["output_dim"] = output_dim_s.str(); |
| (*dict)["sparse_grad"] = sparse_grad_s.str(); |
| (*dict)["dtype"] = MXNetTypeWithBool2String(dtype); |
| } |
| }; |
| |
| /*! |
| * \brief CPU/GPU: Return the amount of temporary storage in bytes required by |
| AddTakeGradLargeBatch |
| * \param num_items number of keys |
| */ |
| template <typename IndexType, typename xpu> |
| inline typename std::enable_if<std::is_same<xpu, cpu>::value, size_t>::type |
| AddTakeGradLargeBatchWorkspaceSize(size_t num_keys) { |
| return 0; |
| } |
| /*! |
| * \brief CPU/GPU: Return the amount of temporary storage in bytes required by |
| AddTakeGradLargeBatch |
| * \param num_items number of keys |
| */ |
| template <typename IndexType, typename xpu> |
| inline typename std::enable_if<std::is_same<xpu, gpu>::value, size_t>::type |
| AddTakeGradLargeBatchWorkspaceSize(size_t num_keys); |
| /*! |
| * \brief CPU/GPU: Gradient accumulate of embedding matrix. |
| dst[sorted[i]] += src[index[i]] |
| Called when the batchsize of src is larger than the featuredim |
| * \param dst destination |
| * \param sorted the sorted indices |
| * \param index original index of the sorted indices |
| * \param src source output |
| * \param workspace (optional) temporary storage |
| */ |
| template <typename IndexType, typename DType> |
| inline void AddTakeGradLargeBatch(mshadow::Tensor<cpu, 2, DType> dst, |
| const mshadow::Tensor<cpu, 1, IndexType>& sorted, |
| const mshadow::Tensor<cpu, 1, IndexType>& index, |
| const mshadow::Tensor<cpu, 2, DType>& src, |
| mshadow::Tensor<cpu, 1, char>* workspace = nullptr) { |
| for (index_t y = 0; y < sorted.size(0); ++y) { |
| dst[sorted[y]] += src[index[y]]; |
| } |
| } |
| template <typename ParamType> |
| inline bool EmbeddingOpShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| using namespace mshadow; |
| const mxnet::TShape& dshape = (*in_attrs)[embedding::kData]; |
| if (!ndim_is_known(dshape)) |
| return false; |
| const ParamType& param = nnvm::get<ParamType>(attrs.parsed); |
| SHAPE_ASSIGN_CHECK(*in_attrs, embedding::kWeight, Shape2(param.input_dim, param.output_dim)); |
| out_attrs->clear(); |
| |
| mxnet::TShape oshape(dshape.ndim() + 1, -1); |
| for (int i = 0; i < dshape.ndim(); ++i) { |
| oshape[i] = dshape[i]; |
| } |
| oshape[dshape.ndim()] = param.output_dim; |
| |
| out_attrs->push_back(oshape); |
| return shape_is_known(oshape); |
| } |
| |
| template <typename ParamType> |
| inline bool EmbeddingOpType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_type, |
| std::vector<int>* out_type) { |
| const ParamType& param = nnvm::get<ParamType>(attrs.parsed); |
| CHECK_EQ(in_type->size(), 2U); |
| CHECK_GE(out_type->size(), 1U); |
| int itype = (*in_type)[0]; |
| CHECK_NE(itype, -1) << "First input must have specified type"; |
| int dtype_in = (*in_type)[1]; |
| int dtype_out = (*out_type)[0]; |
| int dtype = param.dtype; |
| if (dtype_in != -1 && dtype_out != -1) { |
| // Both types defined, make sure they are the same |
| CHECK_EQ(dtype_in, dtype_out) << "Input and output weights must have same type"; |
| dtype = dtype_in; |
| } else if (dtype_in != -1 || dtype_out != -1) { |
| // One of the types defined, choose the one that was defined |
| dtype = (dtype_in != -1) ? dtype_in : dtype_out; |
| } |
| if ((*in_type)[1] == -1) |
| (*in_type)[1] = dtype; |
| out_type->clear(); |
| out_type->push_back(dtype); |
| return true; |
| } |
| |
| // storage type inference function for _backward_Embedding |
| inline bool EmbeddingOpBackwardStorageType(const nnvm::NodeAttrs& attrs, |
| const int dev_mask, |
| DispatchMode* dispatch_mode, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 2U); |
| const bool sparse_grad = nnvm::get<EmbeddingParam>(attrs.parsed).sparse_grad; |
| const NDArrayStorageType target_stype = sparse_grad ? kRowSparseStorage : kDefaultStorage; |
| const auto target_mode = sparse_grad ? DispatchMode::kFComputeEx : DispatchMode::kFCompute; |
| |
| const int ograd_stype = in_attrs->at(0); |
| const int data_stype = in_attrs->at(1); |
| int& data_grad_stype = out_attrs->at(0); |
| int& weight_grad_stype = out_attrs->at(1); |
| bool dispatched = false; |
| if (!dispatched && ograd_stype == kDefaultStorage && data_stype == kDefaultStorage) { |
| // dns, dns -> dns, dns/rsp |
| if (type_assign(&data_grad_stype, kDefaultStorage) && |
| type_assign(&weight_grad_stype, target_stype)) { |
| dispatched = dispatch_mode_assign(dispatch_mode, target_mode); |
| } |
| } |
| // Print user friendly error message to notify misuses of sparse_grad |
| if (weight_grad_stype != target_stype) { |
| LOG(FATAL) << "Cannot use sparse_grad = " << sparse_grad |
| << ", while stype of gradients w.r.t embedding weight is " |
| << common::stype_string(weight_grad_stype); |
| } |
| return dispatched; |
| } |
| |
| /*! \brief TakeNonzeroAxis is designated for general take when |
| * axis is not zero (for CPU optimized version use TakeNonZeroAxisCPU and |
| for axis zero use TakeZeroAxisGPU or TakeZeroAxisCPU) |
| */ |
| template <bool clip = true> |
| struct TakeNonzeroAxis { |
| /*! |
| * \brief Map function for take operator |
| * \param i global thread id |
| * \param out_data ptr to output buffer |
| * \param in_data ptr to input buffer |
| * \param idx ptr to indices buffer |
| * \param in_ndims # of dims of input tensor |
| * \param out_ndims # of dims of output tensor |
| * \param idx_ndims # of dims of indices tensor |
| * \param axis_dim dim size of the axis dimension |
| * \param axis axis id |
| */ |
| template <typename DType, typename IType> |
| MSHADOW_XINLINE static void Map(index_t i, |
| DType* out_data, |
| const DType* in_data, |
| const IType* idx, |
| const index_t out_prev_stride, |
| const index_t in_prev_stride, |
| const index_t in_stride, |
| const int in_ndims, |
| const int out_ndims, |
| const int idx_ndims, |
| const int axis_dim, |
| const int axis) { |
| // i is the global flattened index in the output |
| const index_t out_head_index = i / out_prev_stride; |
| const index_t out_rest_index = i % out_prev_stride; |
| const index_t out_mid_index = out_rest_index / in_stride; |
| const index_t out_tail_index = (axis == in_ndims - 1) ? 0 : (out_rest_index % in_stride); |
| index_t idx_index = static_cast<index_t>(idx[out_mid_index]); |
| if (clip) { |
| idx_index = (idx_index < 0) ? 0 : idx_index; |
| idx_index = (idx_index > axis_dim - 1) ? (axis_dim - 1) : idx_index; |
| } else { |
| idx_index %= axis_dim; |
| idx_index += (idx_index < 0) ? axis_dim : 0; |
| } |
| const index_t in_tail_index = out_tail_index; |
| const index_t in_head_index = out_head_index; |
| index_t in_src_index = in_tail_index + idx_index * in_stride; |
| in_src_index += in_head_index * in_prev_stride; |
| out_data[i] = in_data[in_src_index]; |
| } |
| }; |
| |
| // Embedding forward implementation with dense weight |
| template <typename xpu> |
| void EmbeddingOpForwardDnsImpl(mshadow::Stream<xpu>* s, |
| const TBlob& data, |
| const TBlob& weight, |
| const OpReqType req, |
| const TBlob& output); |
| |
| template <int req> |
| struct TakeRspKernel { |
| /*! |
| * \brief |
| * \param i thread id |
| * \param data input data |
| * \param out output |
| * \param weight_idx indices of rsp weight |
| * \param weight_data data of rsp weight |
| * \param row_length number of elements per row |
| * \param nnr number of non-zero rows |
| */ |
| template <typename DType, typename IType, typename RType> |
| MSHADOW_XINLINE static void Map(index_t i, |
| const IType* data, |
| DType* out, |
| const RType* weight_idx, |
| const DType* weight_data, |
| const nnvm::dim_t row_length, |
| const nnvm::dim_t nnr) { |
| using nnvm::dim_t; |
| const dim_t val = static_cast<dim_t>(data[i]); |
| const DType zero = 0; |
| // Use binary search to find the lower_bound of val in weight_idx array |
| // (adapted based on the binary search in dot kernel) |
| const RType* first = weight_idx; |
| const RType* last = weight_idx + nnr; |
| const RType* it; |
| dim_t count = last - first, step; |
| while (count > 0) { |
| it = first; |
| step = count / 2; |
| it += step; |
| if (*it < val) { |
| first = ++it; |
| count -= step + 1; |
| } else { |
| count = step; |
| } |
| } |
| // end of binary search |
| const dim_t idx_offset = first - weight_idx; |
| const dim_t out_offset = i * row_length; |
| const dim_t weight_offset = idx_offset * row_length; |
| // target idx might be missing in weight.idx. For example, |
| // weight.idx = [5,10] and data = [3,7], so binary search fails to |
| // find any matching indices in weight_idx. |
| if (idx_offset >= nnr || *(weight_idx + idx_offset) > val) { |
| // val not found, fill zeros |
| for (int j = 0; j < row_length; j++) { |
| KERNEL_ASSIGN(out[out_offset + j], req, zero); |
| } |
| } else { |
| for (int j = 0; j < row_length; j++) { |
| KERNEL_ASSIGN(out[out_offset + j], req, weight_data[weight_offset + j]); |
| } |
| } |
| } |
| }; |
| |
| template <typename xpu> |
| inline void EmbeddingOpForwardRspImpl(mshadow::Stream<xpu>* s, |
| const TBlob& data, |
| const NDArray& weight, |
| const OpReqType req, |
| const TBlob& output) { |
| using namespace mxnet_op; |
| using namespace rowsparse; |
| MSHADOW_TYPE_SWITCH(output.type_flag_, DType, { |
| MSHADOW_TYPE_SWITCH(data.type_flag_, IType, { |
| MSHADOW_TYPE_SWITCH(weight.aux_type(kIdx), RType, { |
| MXNET_ASSIGN_REQ_SWITCH(req, req_t, { |
| size_t data_size = data.shape_.Size(); |
| // only using the second dim since weight.ndim() == 2 |
| const nnvm::dim_t row_length = weight.shape()[1]; |
| Kernel<TakeRspKernel<req_t>, xpu>::Launch(s, |
| data_size, |
| data.dptr<IType>(), |
| output.dptr<DType>(), |
| weight.aux_data(kIdx).dptr<RType>(), |
| weight.data().dptr<DType>(), |
| row_length, |
| weight.aux_shape(kIdx)[0]); |
| }); |
| }); |
| }); |
| }); |
| } |
| |
| // Embedding forward implementation with row_sparse weight |
| template <typename xpu> |
| void SparseEmbeddingOpForwardRspImpl(const OpContext& ctx, |
| const TBlob& data, |
| const NDArray& weight, |
| const OpReqType req, |
| const TBlob& output); |
| |
| template <typename xpu> |
| void EmbeddingOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| CHECK_EQ(req[embedding::kOut], kWriteTo); |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 1U); |
| CHECK_EQ(inputs[embedding::kWeight].ndim(), 2U) |
| << "Embedding layer expects its weight to be two-dimensional. " |
| << inputs[embedding::kWeight].ndim() << " dimensional input is given instead"; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| EmbeddingOpForwardDnsImpl<xpu>(s, |
| inputs[embedding::kData], |
| inputs[embedding::kWeight], |
| req[embedding::kOut], |
| outputs[embedding::kOut]); |
| } |
| |
| /*! \brief cast to type and clip to range [0, K - 1] |
| */ |
| struct tcast_clip { |
| template <typename OType, typename IType> |
| MSHADOW_XINLINE static void Map(int i, OType* out_data, const IType* in_data, const OType K) { |
| OType j = static_cast<OType>(in_data[i]); |
| if (j <= 0) |
| j = 0; |
| else if (j >= K) |
| j = K - 1; |
| out_data[i] = j; |
| } |
| }; |
| |
| template <typename xpu> |
| void EmbeddingOpBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 2U); |
| CHECK_EQ(req[embedding::kData], kNullOp) |
| << "Embedding layer doesn't support calculate data gradient"; |
| CHECK_EQ(outputs[1].type_flag_, inputs[0].type_flag_); |
| |
| const mxnet::TShape& ishape = inputs[1].shape_; |
| const mxnet::TShape& oshape = inputs[0].shape_; |
| |
| Stream<xpu>* s = ctx.get_stream<xpu>(); |
| |
| bool safe_acc = dmlc::GetEnv("MXNET_SAFE_ACCUMULATION", true); |
| if (!safe_acc && outputs[1].type_flag_ == mshadow::kFloat16) { |
| common::LogOnce( |
| "MXNET_SAFE_ACCUMULATION=1 is recommended for EmbeddingOpBackward " |
| "with float16 inputs. " |
| "See https://mxnet.apache.org/api/faq/env_var " |
| "for more details."); |
| } |
| MXNET_REAL_ACC_TYPE_SWITCH(outputs[1].type_flag_, DType, AType, { |
| MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { |
| Tensor<xpu, 1, IType> data = |
| inputs[1].get_with_shape<xpu, 1, IType>(Shape1(ishape.ProdShape(0, ishape.ndim())), s); |
| Tensor<xpu, 2, DType> grad_out = inputs[0].get_with_shape<xpu, 2, DType>( |
| Shape2(oshape.ProdShape(0, oshape.ndim() - 1), oshape[oshape.ndim() - 1]), s); |
| Tensor<xpu, 2, DType> grad_in = outputs[1].get<xpu, 2, DType>(s); |
| |
| if (req[embedding::kWeight] == kWriteTo || req[embedding::kWeight] == kAddTo) { |
| if (req[embedding::kWeight] == kWriteTo) { |
| grad_in = scalar<DType>(0.0f); |
| } |
| if (safe_acc) { |
| // Temporary storage for safe accumulation |
| size_t temp_space_size = grad_in.size(0) * grad_in.size(1) * sizeof(AType); |
| Tensor<xpu, 1, char> temp_space = |
| ctx.requested[embedding::kTempSpace].get_space_typed<xpu, 1, char>( |
| Shape1(temp_space_size), s); |
| Tensor<xpu, 2, AType> temp_grad_in( |
| reinterpret_cast<AType*>(temp_space.dptr_), grad_in.shape_, s); |
| AddTakeGrad(grad_in, temp_grad_in, data, grad_out); |
| } else { |
| AddTakeGrad(grad_in, data, grad_out); |
| } |
| } else { |
| LOG(FATAL) << "wrong req"; |
| } |
| }); |
| }); |
| } |
| |
| struct AddTakeGradRspKernel { |
| /*! |
| * \brief Each thread i is responsible for row slices in [segment_start, segment_end) |
| of the result gradient |
| * \param tid global thread id |
| * \param grad the gradient to calculate |
| * \param prefix_sum the inclusive prefix sum of row ids of the gradient |
| * \param ograd output gradient |
| * \param row_length the length of the row slices of the gradient |
| * \param data_val the values of input data |
| * \param data_size number of values of input data |
| * \param segment_length the length of row segment to process for each thread |
| * \param nnr total number of non-zero rows of result gradient |
| */ |
| template <typename DType, typename IType> |
| MSHADOW_CINLINE static void Map(int tid, |
| DType* grad, |
| const nnvm::dim_t* prefix_sum, |
| const DType* ograd, |
| const nnvm::dim_t row_length, |
| const IType* data_val, |
| const nnvm::dim_t data_size, |
| const nnvm::dim_t segment_length, |
| const nnvm::dim_t nnr) { |
| using nnvm::dim_t; |
| dim_t segment_start = tid * segment_length; |
| dim_t segment_end = std::min(nnr, segment_start + segment_length); |
| // scan all data |
| for (dim_t data_i = 0; data_i < data_size; data_i++) { |
| dim_t data = static_cast<dim_t>(data_val[data_i]); |
| dim_t grad_row_id = prefix_sum[data] - 1; |
| if (grad_row_id < segment_start || grad_row_id >= segment_end) |
| continue; |
| // no projection is performed |
| dim_t ograd_i = data_i * row_length; |
| dim_t grad_i = grad_row_id * row_length; |
| for (dim_t offset = 0; offset < row_length; offset++) { |
| grad[grad_i + offset] += ograd[ograd_i + offset]; |
| } |
| } |
| } |
| }; |
| |
| template <typename xpu> |
| inline void SparseEmbeddingOpBackwardRspImpl(const bool deterministic, |
| const OpContext& ctx, |
| const TBlob& ograd, |
| const TBlob& data, |
| const OpReqType req, |
| const NDArray& output); |
| |
| template <typename xpu> |
| void EmbeddingOpBackwardEx(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 2U); |
| const NDArray& weight_grad = outputs[1]; |
| const NDArray& ograd = inputs[0]; |
| const NDArray& data = inputs[1]; |
| // check dtype |
| CHECK_EQ(weight_grad.dtype(), ograd.dtype()); |
| // check req |
| CHECK_EQ(req[embedding::kData], kNullOp) |
| << "Embedding layer doesn't support calculate data gradient"; |
| if (data.storage_type() == kDefaultStorage && ograd.storage_type() == kDefaultStorage && |
| weight_grad.storage_type() == kRowSparseStorage) { |
| SparseEmbeddingOpBackwardRspImpl<xpu>( |
| true, ctx, ograd.data(), data.data(), req[embedding::kWeight], weight_grad); |
| } else { |
| LogUnimplementedOp(attrs, ctx, inputs, req, outputs); |
| } |
| } |
| |
| namespace take_ { // to avoid name conflict |
| enum TakeOpInputs { kArr, kIdx }; |
| enum TakeOpOutputs { kOut }; |
| enum TakeOpResource { kTempSpace }; |
| enum TakeOpMode { kRaise, kWrap, kClip }; |
| } // namespace take_ |
| |
| // TODO(somebody): behaviors specified by params |
| struct TakeParam : public dmlc::Parameter<TakeParam> { |
| int axis; |
| int mode; |
| DMLC_DECLARE_PARAMETER(TakeParam) { |
| DMLC_DECLARE_FIELD(axis).set_default(0).describe( |
| "The axis of input array to be taken." |
| "For input tensor of rank r, it could be in the range of [-r, r-1]"); |
| DMLC_DECLARE_FIELD(mode) |
| .add_enum("raise", take_::kRaise) |
| .add_enum("wrap", take_::kWrap) |
| .add_enum("clip", take_::kClip) |
| .set_default(take_::kClip) |
| .describe( |
| "Specify how out-of-bound indices bahave. Default is \"clip\"." |
| " \"clip\" means clip to the range. So, if all indices mentioned are too large," |
| " they are replaced by the index that addresses the last element along an axis." |
| " \"wrap\" means to wrap around." |
| " \"raise\" means to raise an error when index out of range."); |
| } |
| |
| void SetAttrDict(std::unordered_map<std::string, std::string>* dict) { |
| std::ostringstream axis_s, mode_s; |
| axis_s << axis; |
| mode_s << mode; |
| (*dict)["axis"] = axis_s.str(); |
| (*dict)["mode"] = mode_s.str(); |
| switch (mode) { |
| case take_::kRaise: |
| (*dict)["mode"] = "raise"; |
| break; |
| case take_::kClip: |
| (*dict)["mode"] = "clip"; |
| break; |
| case take_::kWrap: |
| (*dict)["mode"] = "wrap"; |
| break; |
| default: |
| (*dict)["mode"] = mode_s.str(); |
| } |
| } |
| }; |
| |
| inline bool TakeOpShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| using namespace mshadow; |
| const mxnet::TShape& arrshape = (*in_attrs)[take_::kArr]; |
| const mxnet::TShape& idxshape = (*in_attrs)[take_::kIdx]; |
| if (!shape_is_known(idxshape)) |
| return false; |
| const TakeParam& param = nnvm::get<TakeParam>(attrs.parsed); |
| CHECK(param.axis >= -1 * arrshape.ndim() && param.axis < arrshape.ndim()) |
| << "Axis should be in the range of [-r, r-1] where r is the rank of input tensor"; |
| |
| out_attrs->clear(); |
| |
| const index_t actual_axis = param.axis + ((param.axis < 0) ? arrshape.ndim() : 0); |
| mxnet::TShape oshape(idxshape.ndim() + arrshape.ndim() - 1, -1); |
| for (index_t i = 0; i < idxshape.ndim(); ++i) { |
| oshape[i + actual_axis] = idxshape[i]; |
| } |
| for (index_t i = 0; i < arrshape.ndim(); i++) { |
| if (i < actual_axis) { |
| oshape[i] = arrshape[i]; |
| } else if (i > actual_axis) { |
| oshape[i + idxshape.ndim() - 1] = arrshape[i]; |
| } |
| } |
| out_attrs->push_back(oshape); |
| return shape_is_known(oshape); |
| } |
| |
| inline bool TakeOpType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| CHECK_NE((*in_attrs)[1], -1) << "Index type must be set for take operator"; |
| |
| TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]); |
| TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]); |
| return (*in_attrs)[0] != -1; |
| } |
| |
| // storage type inference function for take |
| inline bool TakeOpForwardStorageType(const nnvm::NodeAttrs& attrs, |
| const int dev_mask, |
| DispatchMode* dispatch_mode, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| const int& idx_stype = in_attrs->at(take_::kIdx); |
| const int& arr_stype = in_attrs->at(take_::kArr); |
| int& out_stype = out_attrs->at(take_::kOut); |
| bool dispatched = false; |
| const TakeParam& param = nnvm::get<TakeParam>(attrs.parsed); |
| if (!dispatched && idx_stype == kDefaultStorage && arr_stype == kDefaultStorage) { |
| // dns, dns -> dns |
| dispatched = |
| storage_type_assign(&out_stype, kDefaultStorage, dispatch_mode, DispatchMode::kFCompute); |
| } |
| if (!dispatched && idx_stype == kDefaultStorage && arr_stype == kCSRStorage && param.axis == 0 && |
| (param.mode == take_::kWrap || param.mode == take_::kClip)) { |
| // take(dns, csr, axis=0) -> csr |
| dispatched = |
| storage_type_assign(&out_stype, kCSRStorage, dispatch_mode, DispatchMode::kFComputeEx); |
| } |
| if (!dispatched) { |
| dispatched = dispatch_fallback(out_attrs, dispatch_mode); |
| } |
| return dispatched; |
| } |
| |
| template <typename xpu> |
| void TakeOpForwardCsrImpl(const TakeParam& params, |
| const OpContext& ctx, |
| const TBlob& idx, |
| const NDArray& arr, |
| OpReqType req, |
| const NDArray& output); |
| |
| template <typename xpu> |
| void TakeOpForwardEx(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<NDArray>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<NDArray>& outputs) { |
| CHECK_EQ(req[take_::kOut], kWriteTo); |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 1U); |
| const NDArray& idx = inputs[take_::kIdx]; |
| const NDArray& arr = inputs[take_::kArr]; |
| const NDArray& out = outputs[take_::kOut]; |
| const auto idx_stype = idx.storage_type(); |
| const auto arr_stype = arr.storage_type(); |
| const auto out_stype = out.storage_type(); |
| const auto params = nnvm::get<TakeParam>(attrs.parsed); |
| if (idx_stype == kDefaultStorage && arr_stype == kCSRStorage && out_stype == kCSRStorage) { |
| // dns, csr -> csr |
| TakeOpForwardCsrImpl<xpu>(params, ctx, idx.data(), arr, req[0], out); |
| } else { |
| LogUnimplementedOp(attrs, ctx, inputs, req, outputs); |
| } |
| } |
| |
| template <typename xpu> |
| void TakeOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs); |
| |
| struct TakeGradGeneralKernel { |
| /*! |
| * \brief Map function for general case of take grad |
| * \param tid global thread id |
| * \param arr_grad ptr to in_grad |
| * \param ograd ptr to out_grad |
| * \param src_indptr ptr to indptr to src indices |
| * \param original_idx ptr to original indices of the inputs |
| * \param in_strides strides of inputs |
| * \param out_strides strides of outputs |
| * \param in_ndims # of dims of input tensor |
| * \param out_ndims # of dims of output tensor |
| * \param idx_ndims # of dims of indices tensor |
| * \param axis_dim dim size of the axis dimension |
| * \param axis axis id |
| */ |
| template <typename DType, typename IType> |
| MSHADOW_XINLINE static void Map(int tid, |
| DType* arr_grad, |
| const DType* ograd, |
| const IType* src_indptr, |
| const IType* original_idx, |
| mshadow::Shape<10> in_strides, |
| mshadow::Shape<10> out_strides, |
| const int in_ndims, |
| const int out_ndims, |
| const int idx_ndims, |
| const int axis, |
| const int K) { |
| const int in_head_index = (axis == 0) ? 0 : tid / in_strides[axis - 1]; |
| const int in_rest_index = (axis == 0) ? tid : tid % in_strides[axis - 1]; |
| const int in_mid_index = in_rest_index / in_strides[axis]; |
| const int in_tail_index = (axis == in_ndims - 1) ? 0 : (in_rest_index % in_strides[axis]); |
| for (IType i = src_indptr[in_mid_index]; i < src_indptr[in_mid_index + 1]; ++i) { |
| int out_mid_index = original_idx[i]; |
| out_mid_index = (out_mid_index < 0) ? out_mid_index + K : out_mid_index; |
| int target = in_tail_index + out_mid_index * in_strides[axis]; |
| target += (axis == 0) ? 0 : in_head_index * out_strides[axis - 1]; |
| arr_grad[tid] += ograd[target]; |
| } |
| } |
| }; |
| |
| struct TakeGradGeneralKernelSafeAccumulation { |
| /*! |
| * \brief Map function for general case of take grad |
| * \param tid global thread id |
| * \param arr_grad ptr to in_grad |
| * \param temp ptr to temporal space to perform accumulation |
| * \param ograd ptr to out_grad |
| * \param src_indptr ptr to indptr to src indices |
| * \param original_idx ptr to original indices of the inputs |
| * \param in_strides strides of inputs |
| * \param out_strides strides of outputs |
| * \param in_ndims # of dims of input tensor |
| * \param out_ndims # of dims of output tensor |
| * \param idx_ndims # of dims of indices tensor |
| * \param axis_dim dim size of the axis dimension |
| * \param axis axis id |
| */ |
| template <typename DType, typename IType, typename AType> |
| MSHADOW_XINLINE static void Map(int tid, |
| DType* arr_grad, |
| AType* temp, |
| const DType* ograd, |
| const IType* src_indptr, |
| const IType* original_idx, |
| mshadow::Shape<10> in_strides, |
| mshadow::Shape<10> out_strides, |
| const int in_ndims, |
| const int out_ndims, |
| const int idx_ndims, |
| const int axis, |
| const int K) { |
| const int in_head_index = (axis == 0) ? 0 : tid / in_strides[axis - 1]; |
| const int in_rest_index = (axis == 0) ? tid : tid % in_strides[axis - 1]; |
| const int in_mid_index = in_rest_index / in_strides[axis]; |
| const int in_tail_index = (axis == in_ndims - 1) ? 0 : (in_rest_index % in_strides[axis]); |
| temp[tid] = static_cast<AType>(arr_grad[tid]); |
| for (IType i = src_indptr[in_mid_index]; i < src_indptr[in_mid_index + 1]; ++i) { |
| int out_mid_index = original_idx[i]; |
| out_mid_index = (out_mid_index < 0) ? out_mid_index + K : out_mid_index; |
| int target = in_tail_index + out_mid_index * in_strides[axis]; |
| target += (axis == 0) ? 0 : in_head_index * out_strides[axis - 1]; |
| temp[tid] += ograd[target]; |
| } |
| arr_grad[tid] = temp[tid]; |
| } |
| }; |
| |
| template <bool clip = true, bool safe_acc = false, typename AType> |
| void TakeOpBackwardImpl(mshadow::Stream<cpu>* s, |
| const OpContext& ctx, |
| const TBlob& arr, |
| const TBlob& idx, |
| const TBlob& ograd, |
| const int axis) { |
| using namespace mxnet_op; |
| using namespace mshadow; |
| CHECK(axis != 0) << "axis == 0 case should be dispatched to the legacy implementation"; |
| const mxnet::TShape& arrshape = arr.shape_; |
| const mxnet::TShape& idxshape = idx.shape_; |
| const mxnet::TShape& oshape = ograd.shape_; |
| MSHADOW_TYPE_SWITCH(idx.type_flag_, IType, { |
| // get size of temporary storage for sort |
| int* src_indptr_ptr = nullptr; |
| size_t temp_storage_bytes = SortByKeyWorkspaceSize<int, int, cpu>(idxshape.Size()); |
| size_t original_idx_bytes = idxshape.Size() * sizeof(int); |
| size_t src_indptr_bytes = (arrshape[axis] + 1) * sizeof(int); |
| size_t temp_accumulation_arrgrad_bytes = 0; |
| if (safe_acc) { |
| temp_accumulation_arrgrad_bytes = arr.Size() * sizeof(AType); |
| } |
| size_t workspace_bytes = src_indptr_bytes + 2 * original_idx_bytes + temp_storage_bytes + |
| temp_accumulation_arrgrad_bytes; |
| Tensor<cpu, 1, char> workspace = |
| ctx.requested[0].get_space_typed<cpu, 1, char>(Shape1(workspace_bytes), s); |
| AType* temp_accum_arrgrad_ptr = reinterpret_cast<AType*>(workspace.dptr_); |
| int* sorted_idx_ptr = reinterpret_cast<int*>(workspace.dptr_ + temp_accumulation_arrgrad_bytes); |
| int* original_idx_ptr = reinterpret_cast<int*>(workspace.dptr_ + original_idx_bytes + |
| temp_accumulation_arrgrad_bytes); |
| src_indptr_ptr = reinterpret_cast<int*>(workspace.dptr_ + 2 * original_idx_bytes + |
| temp_accumulation_arrgrad_bytes); |
| Tensor<cpu, 1, char> temp_storage(workspace.dptr_ + 2 * original_idx_bytes + src_indptr_bytes + |
| temp_accumulation_arrgrad_bytes, |
| Shape1(temp_storage_bytes), |
| s); |
| // Reset indptr to zero |
| Kernel<set_zero, cpu>::Launch(s, arrshape[axis] + 1, src_indptr_ptr); |
| // Fill original_idx |
| Kernel<range_fwd, cpu>::Launch(s, idxshape.Size(), 1, 0, 1, kWriteTo, original_idx_ptr); |
| // Fill sorted_idx_ptr with unsorted copy of idx |
| Kernel<mshadow_op::identity_with_cast, cpu>::Launch( |
| s, idxshape.Size(), sorted_idx_ptr, idx.dptr<IType>()); |
| if (clip) { |
| Kernel<op_with_req<mshadow_op::clip, kWriteTo>, cpu>::Launch( |
| s, |
| idxshape.Size(), |
| sorted_idx_ptr, |
| sorted_idx_ptr, |
| 0, |
| static_cast<int>(arrshape[axis] - 1)); |
| } else { |
| Kernel<op_with_req<mshadow_op::mod, kWriteTo>, cpu>::Launch( |
| s, idxshape.Size(), sorted_idx_ptr, sorted_idx_ptr, static_cast<int>(arrshape[axis])); |
| } |
| Tensor<cpu, 1, int> original_idx(original_idx_ptr, Shape1(idxshape.Size()), s); |
| int num_bits = common::ilog2ui(static_cast<unsigned int>(idxshape.Size()) - 1); |
| Tensor<cpu, 1, int> sorted_idx(sorted_idx_ptr, Shape1(idxshape.Size()), s); |
| SortByKey(sorted_idx, original_idx, true, &temp_storage, 0, num_bits); |
| for (size_t i = 0; i < idxshape.Size(); ++i) { |
| src_indptr_ptr[sorted_idx_ptr[i] + 1] += 1; |
| } |
| for (int i = 0; i < arrshape[axis]; ++i) { |
| src_indptr_ptr[i + 1] += src_indptr_ptr[i]; |
| } |
| Shape<10> in_strides; |
| int stride = 1; |
| for (int i = arrshape.ndim() - 1; i >= 0; stride *= arrshape[i], --i) { |
| in_strides[i] = stride; |
| } |
| Shape<10> out_strides; |
| stride = 1; |
| for (int i = oshape.ndim() - 1; i >= 0; stride *= oshape[i], --i) { |
| out_strides[i] = stride; |
| } |
| MSHADOW_TYPE_SWITCH(arr.type_flag_, DType, { |
| if (safe_acc) { |
| Kernel<TakeGradGeneralKernelSafeAccumulation, cpu>::Launch( |
| s, |
| arrshape.Size(), |
| arr.dptr<DType>(), |
| temp_accum_arrgrad_ptr, |
| ograd.dptr<DType>(), |
| src_indptr_ptr, |
| original_idx_ptr, |
| in_strides, |
| out_strides, |
| arrshape.ndim(), |
| oshape.ndim(), |
| idxshape.ndim(), |
| axis, |
| static_cast<int>(arrshape[axis])); |
| } else { |
| Kernel<TakeGradGeneralKernel, cpu>::Launch(s, |
| arrshape.Size(), |
| arr.dptr<DType>(), |
| ograd.dptr<DType>(), |
| src_indptr_ptr, |
| original_idx_ptr, |
| in_strides, |
| out_strides, |
| arrshape.ndim(), |
| oshape.ndim(), |
| idxshape.ndim(), |
| axis, |
| static_cast<int>(arrshape[axis])); |
| } |
| }); |
| }); |
| } |
| |
| #ifdef __CUDACC__ |
| template <bool clip = true, bool safe_acc = false, typename AType> |
| void TakeOpBackwardImpl(mshadow::Stream<gpu>* s, |
| const OpContext& ctx, |
| const TBlob& arr, |
| const TBlob& idx, |
| const TBlob& ograd, |
| const int axis) { |
| using namespace mxnet_op; |
| using namespace mshadow; |
| CHECK(axis != 0) << "axis == 0 case should be dispatched to the legacy implementation"; |
| const mxnet::TShape& arrshape = arr.shape_; |
| const mxnet::TShape& idxshape = idx.shape_; |
| const mxnet::TShape& oshape = ograd.shape_; |
| MSHADOW_TYPE_SWITCH(idx.type_flag_, IType, { |
| // get size of temporary storage for sort |
| char* temp_storage_ptr = nullptr; |
| size_t scan_temp_storage_bytes = 0; |
| int* src_indptr_ptr = nullptr; |
| cub::DeviceScan::ExclusiveSum(temp_storage_ptr, |
| scan_temp_storage_bytes, |
| src_indptr_ptr, |
| src_indptr_ptr, |
| arrshape[axis] + 1, |
| mshadow::Stream<gpu>::GetStream(s)); |
| size_t sort_temp_storage_bytes = SortByKeyWorkspaceSize<int, int, gpu>(idxshape.Size()); |
| size_t histo_temp_storage_bytes = 0; |
| int* sorted_idx_ptr = nullptr; |
| cub::DeviceHistogram::HistogramEven(temp_storage_ptr, |
| histo_temp_storage_bytes, |
| sorted_idx_ptr, |
| src_indptr_ptr, |
| static_cast<int>(arrshape[axis] + 1), |
| 0, |
| static_cast<int>(arrshape[axis] + 1), |
| static_cast<int>(idxshape.Size()), |
| mshadow::Stream<gpu>::GetStream(s)); |
| size_t temp_storage_bytes = max(scan_temp_storage_bytes, sort_temp_storage_bytes); |
| temp_storage_bytes = max(temp_storage_bytes, histo_temp_storage_bytes); |
| size_t original_idx_bytes = idxshape.Size() * sizeof(int); |
| size_t src_indptr_bytes = (arrshape[axis] + 1) * sizeof(int); |
| size_t temp_accumulation_igrad_bytes = 0; |
| if (safe_acc) { |
| temp_accumulation_igrad_bytes = arr.Size() * sizeof(AType); |
| } |
| size_t workspace_bytes = src_indptr_bytes + 2 * original_idx_bytes + temp_storage_bytes + |
| temp_accumulation_igrad_bytes; |
| Tensor<gpu, 1, char> workspace = |
| ctx.requested[0].get_space_typed<gpu, 1, char>(Shape1(workspace_bytes), s); |
| AType* temp_accum_igrad_ptr = reinterpret_cast<AType*>(workspace.dptr_); |
| sorted_idx_ptr = reinterpret_cast<int*>(workspace.dptr_ + temp_accumulation_igrad_bytes); |
| int* original_idx_ptr = reinterpret_cast<int*>(workspace.dptr_ + original_idx_bytes + |
| temp_accumulation_igrad_bytes); |
| src_indptr_ptr = reinterpret_cast<int*>(workspace.dptr_ + 2 * original_idx_bytes + |
| temp_accumulation_igrad_bytes); |
| temp_storage_ptr = |
| workspace.dptr_ + 2 * original_idx_bytes + src_indptr_bytes + temp_accumulation_igrad_bytes; |
| |
| // Reset indptr to zero |
| Kernel<set_zero, gpu>::Launch(s, arrshape[axis] + 1, src_indptr_ptr); |
| // Fill original_idx |
| Kernel<range_fwd, gpu>::Launch(s, |
| idxshape.Size(), |
| 1, |
| static_cast<int>(0), |
| static_cast<int>(1), |
| kWriteTo, |
| original_idx_ptr); |
| // Fill sorted_idx_ptr with unsorted copy of idx |
| Kernel<mshadow_op::identity_with_cast, gpu>::Launch( |
| s, idxshape.Size(), sorted_idx_ptr, idx.dptr<IType>()); |
| if (clip) { |
| Kernel<op_with_req<mshadow_op::clip, kWriteTo>, gpu>::Launch( |
| s, idxshape.Size(), sorted_idx_ptr, sorted_idx_ptr, 0, static_cast<int>(arrshape[axis])); |
| } else { |
| Kernel<op_with_req<mshadow_op::mod, kWriteTo>, gpu>::Launch( |
| s, idxshape.Size(), sorted_idx_ptr, sorted_idx_ptr, static_cast<int>(arrshape[axis])); |
| } |
| Tensor<gpu, 1, int> original_idx(original_idx_ptr, Shape1(idxshape.Size()), s); |
| Tensor<gpu, 1, char> temp_storage(temp_storage_ptr, Shape1(temp_storage_bytes), s); |
| int num_bits = common::ilog2ui(static_cast<unsigned int>(idxshape.Size()) - 1); |
| Tensor<gpu, 1, int> sorted_idx(sorted_idx_ptr, Shape1(idxshape.Size()), s); |
| SortByKey(sorted_idx, original_idx, true, &temp_storage, 0, num_bits); |
| cub::DeviceHistogram::HistogramEven(temp_storage_ptr, |
| temp_storage_bytes, |
| sorted_idx_ptr, |
| src_indptr_ptr, |
| static_cast<int>(arrshape[axis] + 1), |
| 0, |
| static_cast<int>(arrshape[axis] + 1), |
| static_cast<int>(idxshape.Size()), |
| mshadow::Stream<gpu>::GetStream(s)); |
| cub::DeviceScan::ExclusiveSum(temp_storage_ptr, |
| temp_storage_bytes, |
| src_indptr_ptr, |
| src_indptr_ptr, |
| arrshape[axis] + 1, |
| mshadow::Stream<gpu>::GetStream(s)); |
| |
| Shape<10> in_strides; |
| int stride = 1; |
| for (int i = arrshape.ndim() - 1; i >= 0; stride *= arrshape[i], --i) { |
| in_strides[i] = stride; |
| } |
| Shape<10> out_strides; |
| stride = 1; |
| for (int i = oshape.ndim() - 1; i >= 0; stride *= oshape[i], --i) { |
| out_strides[i] = stride; |
| } |
| MSHADOW_TYPE_SWITCH(arr.type_flag_, DType, { |
| if (safe_acc) { |
| Kernel<TakeGradGeneralKernelSafeAccumulation, gpu>::Launch( |
| s, |
| arrshape.Size(), |
| arr.dptr<DType>(), |
| temp_accum_igrad_ptr, |
| ograd.dptr<DType>(), |
| src_indptr_ptr, |
| original_idx_ptr, |
| in_strides, |
| out_strides, |
| arrshape.ndim(), |
| oshape.ndim(), |
| idxshape.ndim(), |
| axis, |
| static_cast<int>(arrshape[axis])); |
| } else { |
| Kernel<TakeGradGeneralKernel, gpu>::Launch(s, |
| arrshape.Size(), |
| arr.dptr<DType>(), |
| ograd.dptr<DType>(), |
| src_indptr_ptr, |
| original_idx_ptr, |
| in_strides, |
| out_strides, |
| arrshape.ndim(), |
| oshape.ndim(), |
| idxshape.ndim(), |
| axis, |
| static_cast<int>(arrshape[axis])); |
| } |
| }); |
| }); |
| } |
| #endif |
| |
| template <typename xpu> |
| void TakeOpBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using namespace mshadow::expr; |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 2U); |
| CHECK_NE(req[take_::kIdx], kAddTo) |
| << "take layer doesn't support gradient of req type kAddTo to index"; |
| |
| const TakeParam& param = nnvm::get<TakeParam>(attrs.parsed); |
| |
| // grad_out is the gradient of the outputs in the feed-forward |
| // grad_in is the gradient of the inputs in the feed-forward |
| Stream<xpu>* s = ctx.get_stream<xpu>(); |
| |
| bool safe_acc = dmlc::GetEnv("MXNET_SAFE_ACCUMULATION", true); |
| if (!safe_acc && outputs[0].type_flag_ == mshadow::kFloat16) { |
| common::LogOnce( |
| "MXNET_SAFE_ACCUMULATION=1 is recommended for TakeOpBackward " |
| "with float16 inputs. " |
| "See https://mxnet.apache.org/api/faq/env_var " |
| "for more details."); |
| } |
| MXNET_REAL_ACC_TYPE_SWITCH(outputs[0].type_flag_, DType, AType, { |
| MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // index data type |
| // inputs are specified in the .cc file, which are the gradients from |
| // the upper layer and the input index |
| // outputs are the gradients of inputs in the feed-forward pass |
| const mxnet::TShape& idxshape = inputs[1].shape_; |
| const mxnet::TShape& arrshape = outputs[0].shape_; |
| const mxnet::TShape& oshape = inputs[0].shape_; |
| |
| Tensor<xpu, 2, DType> grad_in = outputs[0].get_with_shape<xpu, 2, DType>( |
| Shape2(arrshape[0], arrshape.ProdShape(1, arrshape.ndim())), s); |
| |
| if (req[take_::kArr] == kWriteTo) { |
| grad_in = scalar<DType>(0.0f); |
| } |
| |
| if (idxshape.Size() == 0) { |
| return; |
| } |
| |
| if (req[take_::kIdx] != kNullOp) { |
| mxnet_op::Kernel<mxnet_op::set_zero, xpu>::Launch( |
| s, idxshape.Size(), outputs[take_::kIdx].dptr<IType>()); |
| } |
| |
| const int actual_axis = param.axis + ((param.axis < 0) ? arrshape.ndim() : 0); |
| |
| int idxndim = idxshape.ndim(); |
| Tensor<xpu, 1, IType> idx = |
| inputs[1].get_with_shape<xpu, 1, IType>(Shape1(idxshape.ProdShape(0, idxndim)), s); |
| Tensor<xpu, 2, DType> grad_out = inputs[0].get_with_shape<xpu, 2, DType>( |
| Shape2(oshape.ProdShape(0, idxndim), oshape.ProdShape(idxndim, oshape.ndim())), s); |
| |
| // re-using the previous code for axis = 0 case |
| if (actual_axis == 0) { |
| if (req[take_::kArr] == kWriteTo || req[take_::kArr] == kAddTo) { |
| if (safe_acc) { |
| // Temporary storage for safe accumulation |
| size_t temp_space_size = grad_in.size(0) * grad_in.size(1) * sizeof(AType); |
| Tensor<xpu, 1, char> temp_space = |
| ctx.requested[take_::kTempSpace].get_space_typed<xpu, 1, char>( |
| Shape1(temp_space_size), s); |
| Tensor<xpu, 2, AType> temp_grad_in( |
| reinterpret_cast<AType*>(temp_space.dptr_), grad_in.shape_, s); |
| if (param.mode == take_::kClip) { |
| AddTakeGrad(grad_in, temp_grad_in, idx, grad_out); |
| } else { |
| AddTakeGrad<false>(grad_in, temp_grad_in, idx, grad_out); |
| } |
| } else { |
| if (param.mode == take_::kClip) { |
| AddTakeGrad(grad_in, idx, grad_out); |
| } else { |
| AddTakeGrad<false>(grad_in, idx, grad_out); |
| } |
| } |
| } else { |
| LOG(FATAL) << "wrong req"; |
| } |
| // for all other cases |
| } else { |
| const TBlob& idx = inputs[1]; |
| const TBlob& arr = outputs[0]; |
| const TBlob& ograd = inputs[0]; |
| |
| if (safe_acc) { |
| if (param.mode == take_::kClip) { |
| TakeOpBackwardImpl<true, true, AType>(s, ctx, arr, idx, ograd, actual_axis); |
| } else { |
| TakeOpBackwardImpl<false, true, AType>(s, ctx, arr, idx, ograd, actual_axis); |
| } |
| } else { |
| if (param.mode == take_::kClip) { |
| TakeOpBackwardImpl<true, false, AType>(s, ctx, arr, idx, ograd, actual_axis); |
| } else { |
| TakeOpBackwardImpl<false, false, AType>(s, ctx, arr, idx, ograd, actual_axis); |
| } |
| } |
| } |
| }); |
| }); |
| } |
| |
| inline bool BatchTakeOpShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| LOG(INFO) << "batch_take is deprecated. Please use pick instead."; |
| CHECK_EQ(in_attrs->size(), 2U) << "BatchTake op requires two inputs"; |
| if ((*in_attrs)[1].ndim() != 0) { |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[1]); |
| } else if ((*out_attrs)[0].ndim() != 0) { |
| SHAPE_ASSIGN_CHECK(*in_attrs, 1, (*out_attrs)[0]); |
| } |
| if ((*in_attrs)[0].ndim() == 0) |
| return false; |
| CHECK_GE((*in_attrs)[0].ndim(), 2) << "Data array must have at least 2 dimensional"; |
| if ((*out_attrs)[0].ndim() == 0) |
| return false; |
| CHECK_EQ((*in_attrs)[0].Size() / (*in_attrs)[0][(*in_attrs)[0].ndim() - 1], |
| (*out_attrs)[0].Size()) |
| << "Index array's size must be the same as data array's size excluding the first dimension"; |
| return true; |
| } |
| |
| inline bool BatchTakeOpType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| if ((*in_attrs)[0] != -1) { |
| TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]); |
| } else if ((*out_attrs)[0] != -1) { |
| TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]); |
| } |
| TYPE_ASSIGN_CHECK(*in_attrs, 1, mshadow::kInt32); |
| return true; |
| } |
| |
| /*! \brief take scalar value from 2d data array */ |
| template <int req> |
| struct batch_take { |
| template <typename DType> |
| MSHADOW_XINLINE static void Map(int i, DType* out, const DType* a, const int* idx, int M) { |
| int j = idx[i]; |
| if (j < 0) |
| j = 0; |
| else if (j >= M) |
| j = M - 1; |
| KERNEL_ASSIGN(out[i], req, a[i * M + j]); |
| } |
| }; |
| |
| template <typename xpu> |
| void BatchTakeOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 1U); |
| CHECK_EQ(req.size(), 1U); |
| using namespace mxnet_op; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { |
| MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, { |
| Kernel<batch_take<req_type>, xpu>::Launch(s, |
| outputs[0].Size(), |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<int>(), |
| inputs[0].Size() / inputs[0].shape_[0]); |
| }); |
| }); |
| } |
| |
| /*! |
| * \brief The parameters of the one_hot operator. |
| */ |
| struct OneHotParam : public dmlc::Parameter<OneHotParam> { |
| index_t depth; |
| double on_value; |
| double off_value; |
| int axis; |
| int dtype; |
| DMLC_DECLARE_PARAMETER(OneHotParam) { |
| DMLC_DECLARE_FIELD(depth).describe("Depth of the one hot dimension."); |
| DMLC_DECLARE_FIELD(on_value).set_default(1.0f).describe( |
| "The value assigned to the locations represented by indices."); |
| DMLC_DECLARE_FIELD(off_value).set_default(0.0f).describe( |
| "The value assigned to the locations not represented by indices."); |
| DMLC_DECLARE_FIELD(dtype).set_default(mshadow::kFloat32) |
| MXNET_ADD_ALL_TYPES.describe("DType of the output"); |
| } |
| void SetAttrDict(std::unordered_map<std::string, std::string>* dict) { |
| std::ostringstream depth_s, on_value_s, off_value_s, axis_s, dtype_s; |
| depth_s << depth; |
| on_value_s << on_value; |
| off_value_s << off_value; |
| dtype_s << dtype; |
| (*dict)["depth"] = depth_s.str(); |
| (*dict)["on_value"] = on_value_s.str(); |
| (*dict)["off_value"] = off_value_s.str(); |
| (*dict)["dtype"] = MXNetTypeWithBool2String(dtype); |
| } |
| }; |
| |
| inline void GetOneHotParams(const OneHotParam& param, |
| index_t* depth, |
| double* on_value, |
| double* off_value, |
| int* dtype) { |
| *depth = param.depth; |
| CHECK_GE(*depth, 0) << "Dimension size, depth, must be a non-negative integer"; |
| *on_value = param.on_value; |
| *off_value = param.off_value; |
| *dtype = param.dtype; |
| } |
| |
| inline bool OneHotOpShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| const OneHotParam& param = nnvm::get<OneHotParam>(attrs.parsed); |
| CHECK_EQ(in_attrs->size(), 1U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| // The shape of indices |
| const mxnet::TShape& ishape = (*in_attrs)[0]; |
| if (!shape_is_known(ishape)) |
| return false; |
| |
| index_t depth = 0; |
| double on_value = 1.0; |
| double off_value = 0.0; |
| int dtype = mshadow::kFloat32; |
| GetOneHotParams(param, &depth, &on_value, &off_value, &dtype); |
| |
| mxnet::TShape oshape(ishape.ndim() + 1, -1); |
| for (index_t i = 0; i < ishape.ndim(); ++i) { |
| oshape[i] = ishape[i]; |
| } |
| oshape[oshape.ndim() - 1] = depth; |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, oshape); |
| return shape_is_known(oshape); |
| } |
| |
| inline bool OneHotOpType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 1U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| CHECK_NE((*in_attrs)[0], -1) << "Index type must be set for one_hot operator"; |
| index_t depth = 0; |
| double on_value = 1.0; |
| double off_value = 0.0; |
| int dtype = -1; |
| const OneHotParam& param = nnvm::get<OneHotParam>(attrs.parsed); |
| GetOneHotParams(param, &depth, &on_value, &off_value, &dtype); |
| TYPE_ASSIGN_CHECK(*out_attrs, 0, dtype); // assign output type |
| |
| return true; |
| } |
| |
| template <int req> |
| struct one_hot { |
| template <typename DType, typename IType> |
| MSHADOW_XINLINE static void Map(index_t i, |
| DType* out, |
| const IType* indices, |
| index_t depth, |
| DType on_value) { |
| index_t offset = i * depth; |
| index_t j = static_cast<index_t>(indices[i]); |
| if (j >= 0 && j < depth) { |
| KERNEL_ASSIGN(out[offset + j], req, on_value); |
| } |
| } |
| }; |
| |
| template <typename xpu> |
| void OneHotOpForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| CHECK_EQ(inputs.size(), 1U); |
| CHECK_EQ(outputs.size(), 1U); |
| CHECK_EQ(req.size(), 1U); |
| // The following line is needed to guard the situation when |
| // an output array is empty on GPU. In that case, out.dptr() = 0x0 |
| if (outputs[0].Size() == 0) |
| return; |
| index_t depth = 0; |
| double on_value = 1.0; |
| double off_value = 0.0; |
| int dtype = mshadow::kFloat32; |
| const OneHotParam& param = nnvm::get<OneHotParam>(attrs.parsed); |
| GetOneHotParams(param, &depth, &on_value, &off_value, &dtype); |
| using namespace mxnet_op; |
| using namespace mshadow::expr; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, { // output data type switch |
| mshadow::Tensor<xpu, 1, DType> out = outputs[0].FlatTo1D<xpu, DType>(s); |
| ASSIGN_DISPATCH(out, req[0], static_cast<DType>(off_value)); |
| MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, { // request type switch |
| MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, IType, { // indices data type switch |
| Kernel<one_hot<req_type>, xpu>::Launch(s, |
| inputs[0].Size(), |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<IType>(), |
| depth, |
| static_cast<DType>(on_value)); |
| }); |
| }); |
| }); |
| } |
| |
| struct gather_nd { |
| template <typename DType, typename IType> |
| MSHADOW_XINLINE static void Map(index_t i, |
| OpReqType req, |
| index_t N, |
| index_t M, |
| index_t K, |
| const mshadow::Shape<10> strides, |
| const mshadow::Shape<10> mshape, |
| DType* out, |
| const DType* data, |
| const IType* indices) { |
| index_t offset = 0; |
| for (index_t j = 0; j < M; ++j) { |
| offset += strides[j] * (static_cast<index_t>(indices[j * N + i] + mshape[j]) % mshape[j]); |
| } |
| for (index_t j = 0; j < K; ++j) { |
| KERNEL_ASSIGN(out[i * K + j], req, data[offset + j]); |
| } |
| } |
| }; |
| |
| /*! |
| * \brief If any index in a dimension is out of bound, |
| then the value in this dimension will be set to be the out-of-bound index |
| */ |
| struct is_valid_check_gather_nd { |
| template <typename DType> |
| MSHADOW_XINLINE static void Map(int i, |
| DType* is_valid_dim_ptr, |
| const DType* idx_ptr, |
| const index_t N, |
| const mshadow::Shape<10> mshape) { |
| index_t n = N - 1; |
| while (n >= 0) { |
| if (idx_ptr[i * N + n] < -mshape[i] || idx_ptr[i * N + n] > mshape[i] - 1) { |
| is_valid_dim_ptr[i] = idx_ptr[i * N + n]; |
| break; |
| } |
| n--; |
| } |
| } |
| }; |
| |
| inline bool GatherNDShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| // The shape of indices |
| const mxnet::TShape& dshape = (*in_attrs)[0]; |
| const mxnet::TShape& ishape = (*in_attrs)[1]; |
| |
| if (shape_is_none(dshape) || shape_is_none(ishape)) |
| return false; |
| |
| CHECK_GT(ishape.ndim(), 1) << "gather_nd requires index tensor to have at least 2 dimensions"; |
| |
| CHECK_LE(ishape[0], dshape.ndim()) << "Number of indices exceeds data dimension"; |
| |
| CHECK_LE(ishape[0], 10) << "gather_nd supports indexing along at most 10 dimensions."; |
| |
| mxnet::TShape oshape(ishape.ndim() - 1 + dshape.ndim() - ishape[0], -1); |
| |
| for (int i = 0; i < ishape.ndim() - 1; ++i) { |
| oshape[i] = ishape[i + 1]; |
| } |
| for (int i = 0; i < dshape.ndim() - ishape[0]; ++i) { |
| oshape[ishape.ndim() - 1 + i] = dshape[ishape[0] + i]; |
| } |
| |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, oshape); |
| return shape_is_known(oshape); |
| } |
| |
| inline bool GatherNDType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]); |
| TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]); |
| return true; |
| } |
| |
| struct ScatterNDParam : public dmlc::Parameter<ScatterNDParam> { |
| mxnet::TShape shape; |
| DMLC_DECLARE_PARAMETER(ScatterNDParam) { |
| DMLC_DECLARE_FIELD(shape).describe("Shape of output."); |
| } |
| }; |
| |
| inline bool ScatterNDShape(const nnvm::NodeAttrs& attrs, |
| mxnet::ShapeVector* in_attrs, |
| mxnet::ShapeVector* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| const auto& params = dmlc::get<ScatterNDParam>(attrs.parsed); |
| |
| SHAPE_ASSIGN_CHECK(*out_attrs, 0, params.shape); |
| |
| const mxnet::TShape& dshape = (*in_attrs)[0]; |
| const mxnet::TShape& ishape = (*in_attrs)[1]; |
| const mxnet::TShape& oshape = (*out_attrs)[0]; |
| |
| if (shape_is_none(dshape) || shape_is_none(ishape) || shape_is_none(oshape)) |
| return false; |
| |
| CHECK_GT(ishape.ndim(), 1) << "scatter_nd requires index tensor to have at least 2 dimensions"; |
| |
| CHECK_LE(ishape[0], oshape.ndim()) |
| << "Number of indices exceeds output dimension in operator scatter_nd"; |
| |
| CHECK_LE(ishape[0], 10) << "scatter_nd supports indexing along at most 10 dimensions."; |
| |
| bool valid = dshape.ndim() == ishape.ndim() - 1 + oshape.ndim() - ishape[0]; |
| |
| for (int i = 0; i < ishape.ndim() - 1; ++i) { |
| valid = valid && dshape[i] == ishape[i + 1]; |
| } |
| for (int i = 0; i < oshape.ndim() - ishape[0]; ++i) { |
| valid = valid && dshape[ishape.ndim() - 1 + i] == oshape[ishape[0] + i]; |
| } |
| |
| CHECK(valid) << "Invalid data, indices, and output shape combination for scatter_nd: " << dshape |
| << ", " << ishape << ", " << oshape; |
| |
| return true; |
| } |
| |
| inline bool ScatterNDType(const nnvm::NodeAttrs& attrs, |
| std::vector<int>* in_attrs, |
| std::vector<int>* out_attrs) { |
| CHECK_EQ(in_attrs->size(), 2U); |
| CHECK_EQ(out_attrs->size(), 1U); |
| TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]); |
| TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[0]); |
| return in_attrs->at(0) != -1 && in_attrs->at(1) != -1; |
| } |
| |
| struct scatter_nd { |
| template <typename DType, typename IType> |
| MSHADOW_XINLINE static void Map(index_t i, |
| OpReqType req, |
| index_t N, |
| index_t M, |
| index_t K, |
| const mshadow::Shape<10> strides, |
| DType* out, |
| const DType* data, |
| const IType* indices) { |
| index_t offset = 0; |
| for (index_t j = 0; j < M; ++j) { |
| offset += strides[j] * static_cast<index_t>(indices[j * N + i]); |
| } |
| for (index_t j = 0; j < K; ++j) { |
| KERNEL_ASSIGN(out[offset + j], req, data[i * K + j]); |
| } |
| } |
| }; |
| |
| template <typename xpu> |
| void ScatterNDForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using nnvm::dim_t; |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 1U); |
| if (req[0] == kNullOp) |
| return; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| const mxnet::TShape& oshape = outputs[0].shape_; |
| const mxnet::TShape& ishape = inputs[1].shape_; |
| dim_t M = ishape[0]; |
| dim_t N = ishape.Size() / M; |
| dim_t K = oshape.ProdShape(M, oshape.ndim()); |
| mshadow::Shape<10> strides; |
| for (dim_t i = M - 1, stride = K; i >= 0; stride *= oshape[i], --i) |
| strides[i] = stride; |
| if (kWriteTo == req[0]) { |
| Fill<true>(s, outputs[0], req[0], 0); |
| } |
| MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[0].type_flag_, DType, { // output data type switch |
| MSHADOW_TYPE_SWITCH_WITH_BOOL(inputs[1].type_flag_, IType, { // indices data type switch |
| mxnet_op::Kernel<scatter_nd, xpu>::Launch(s, |
| N, |
| req[0], |
| N, |
| M, |
| K, |
| strides, |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<IType>()); |
| }); |
| }); |
| } |
| |
| template <typename DType, typename IType> |
| inline typename std::enable_if<(!std::is_same<DType, mshadow::half::half_t>::value), void>::type |
| GatherNDBackwardImpl(index_t N, |
| index_t M, |
| index_t K, |
| const mshadow::Shape<10> strides, |
| DType* out, |
| const DType* data, |
| const IType* indices, |
| mshadow::Stream<cpu>* s); |
| |
| template <typename DType, typename IType> |
| inline typename std::enable_if<std::is_same<DType, mshadow::half::half_t>::value, void>::type |
| GatherNDBackwardImpl(index_t N, |
| index_t M, |
| index_t K, |
| const mshadow::Shape<10> strides, |
| DType* out, |
| const DType* data, |
| const IType* indices, |
| mshadow::Stream<cpu>* s); |
| |
| template <typename DType, typename IType> |
| inline void GatherNDBackwardImpl(index_t N, |
| index_t M, |
| index_t K, |
| const mshadow::Shape<10> strides, |
| DType* out, |
| const DType* data, |
| const IType* indices, |
| mshadow::Stream<gpu>* s); |
| |
| template <typename xpu> |
| void GatherNDBackward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| using namespace mshadow; |
| using nnvm::dim_t; |
| CHECK_EQ(inputs.size(), 2U); |
| CHECK_EQ(outputs.size(), 1U); |
| if (req[0] == kNullOp) |
| return; |
| mshadow::Stream<xpu>* s = ctx.get_stream<xpu>(); |
| const mxnet::TShape& oshape = outputs[0].shape_; |
| const mxnet::TShape& ishape = inputs[1].shape_; |
| dim_t M = ishape[0]; |
| dim_t N = ishape.Size() / M; |
| dim_t K = oshape.ProdShape(M, oshape.ndim()); |
| mshadow::Shape<10> strides; |
| for (dim_t i = M - 1, stride = K; i >= 0; stride *= oshape[i], --i) |
| strides[i] = stride; |
| if (kWriteTo == req[0]) { |
| Fill<true>(s, outputs[0], req[0], 0); |
| } |
| MXNET_NO_INT8_TYPE_SWITCH(inputs[0].type_flag_, DType, { // output data type switch |
| MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // indices data type switch |
| GatherNDBackwardImpl(N, |
| M, |
| K, |
| strides, |
| outputs[0].dptr<DType>(), |
| inputs[0].dptr<DType>(), |
| inputs[1].dptr<IType>(), |
| s); |
| }); |
| }); |
| } |
| |
| /*! |
| * This is for internal use only. |
| * DO NOT call this function unless you have to. |
| */ |
| template <typename xpu> |
| void ScatterSetNDForward(const nnvm::NodeAttrs& attrs, |
| const OpContext& ctx, |
| const std::vector<TBlob>& inputs, |
| const std::vector<OpReqType>& req, |
| const std::vector<TBlob>& outputs) { |
| CHECK_EQ(inputs.size(), 3U); |
| CHECK_EQ(outputs.size(), 1U); |
| CHECK_EQ(inputs[0].dptr_, outputs[0].dptr_); |
| ScatterNDForward<xpu>(attrs, ctx, {inputs[1], inputs[2]}, {kWriteInplace}, outputs); |
| } |
| |
| } // namespace op |
| } // namespace mxnet |
| |
| #endif // MXNET_OPERATOR_TENSOR_INDEXING_OP_H_ |
