[new API] add paddle.kthvalue and paddle.Tensor.kthvalue #38386
Conversation
|
Thanks for your contribution! |
|
|
||
| template <typename T> | ||
| bool SortKthvalue(const platform::CUDADeviceContext& ctx, | ||
| const framework::Tensor* input_tensor, const int64_t num_cols, |
There was a problem hiding this comment.
Too many duplicate codes with top_k_function_cuda.h.
| indices)) { | ||
| return; | ||
| } else { | ||
| LOG(INFO) << "KthvalueOP: Some errors happened when use cub sorting"; |
There was a problem hiding this comment.
Throw exception if there is any error. Do not just logging.
| } | ||
| return; | ||
| } else { | ||
| LOG(INFO) << "KthvalueOP: Some errors happened when use cub sorting"; |
| const int& dim_size = input_dims.size(); | ||
| int axis = static_cast<int>(ctx->Attrs().Get<int>("axis")); | ||
| PADDLE_ENFORCE_EQ( | ||
| (axis < dim_size) && (axis >= (-1 * dim_size)), true, |
There was a problem hiding this comment.
Try to write:
PADDLE_ENFORCE_LT(axis, dim_size)
and
PADDLE_ENFORCE_GE(axis, -dim_size)
|
|
||
| protected: | ||
| framework::OpKernelType GetExpectedKernelType( | ||
| const framework::ExecutionContext& ctx) const override { |
There was a problem hiding this comment.
Just write:
return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(ctx, "X"), ctx.GetPlace())
| class KthvalueCPUKernel : public framework::OpKernel<T> { | ||
| public: | ||
| void Compute(const framework::ExecutionContext& context) const override { | ||
| auto* input = context.Input<framework::Tensor>("X"); |
There was a problem hiding this comment.
Too many duplicate codes in CPU/CUDA kernels.
There was a problem hiding this comment.
LGTM.
TODO: do code clean to avoid too many duplicate codes. @zoooo0820
PR types
New features
PR changes
OPs
Describe
This pr adds API/OPs paddle.kthvalue and paddle.Tensor.kthvalue.