random selected rows value

paddle/fluid/operators/uniform_random_op.cc

@@ -24,7 +24,15 @@ template <typename T>
 class CPUUniformRandomKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    auto* tensor = ctx.Output<framework::Tensor>("Out");
+    framework::Tensor* tensor(nullptr);
+    auto out_var = ctx.OutputVar("Out");
+    if (out_var->IsType<framework::LoDTensor>()) {
+      tensor = ctx.Output<framework::LoDTensor>("Out");
+    } else if (out_var->IsType<framework::SelectedRows>()) {
+      tensor = ctx.Output<framework::SelectedRows>("Out")->mutable_value();
+    } else {
+      PADDLE_THROW("Only support LoDTensor and SelectedRows.");
+    }
     T* data = tensor->mutable_data<T>(ctx.GetPlace());
     unsigned int seed = static_cast<unsigned int>(ctx.Attr<int>("seed"));
     std::minstd_rand engine;
@@ -36,6 +44,7 @@ class CPUUniformRandomKernel : public framework::OpKernel<T> {
         static_cast<T>(ctx.Attr<float>("min")),
         static_cast<T>(ctx.Attr<float>("max")));
     int64_t size = tensor->numel();
+    VLOG(3) << "size = " << size;
     for (int64_t i = 0; i < size; ++i) {
       data[i] = dist(engine);
     }
@@ -55,6 +64,7 @@ class UniformRandomOp : public framework::OperatorWithKernel {
         "uniform_random's min must less then max");
     auto& shape = ctx->Attrs().Get<std::vector<int>>("shape");
     std::vector<int64_t> temp;
+    VLOG(3) << "shape.size() = " << shape.size();
     temp.reserve(shape.size());
     for (auto dim : shape) {
       temp.push_back(static_cast<int64_t>(dim));