percentile

apache · Feb 17, 2020 · 0b6a974 · 0b6a974
1 parent afcf4f7
commit 0b6a974
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/operator/numpy/np_percentile_op.cc b/src/operator/numpy/np_percentile_op.cc
@@ -93,7 +93,11 @@ NNVM_REGISTER_OP(_npi_percentile)
 .set_attr<nnvm::FInferType>("FInferType", NumpyPercentileType)
 .set_attr<nnvm::FListInputNames>("FListInputNames",
   [](const NodeAttrs& attrs) {
-    return std::vector<std::string>{"a", "q"};
+     const NumpyPercentileParam& param =
+     nnvm::get<NumpyPercentileParam>(attrs.parsed);
+     return param.q_scalar.has_value() ?
+            std::vector<std::string>{"a"} :
+            std::vector<std::string>{"a", "q"};
   })
 .set_attr<FCompute>("FCompute<cpu>", NumpyPercentileForward<cpu>)
 .set_attr<FResourceRequest>("FResourceRequest",