/*
 * 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.
 */

/*!
 *  Copyright (c) 2016 by Contributors
 * \file sample_op.h
 * \brief Elementary sampling operators
 */
#ifndef MXNET_OPERATOR_RANDOM_SAMPLE_OP_H_
#define MXNET_OPERATOR_RANDOM_SAMPLE_OP_H_

#include <mxnet/operator_util.h>
#include <mshadow/base.h>
#include <string>
#include <vector>
#include "../../common/utils.h"
#include "../mxnet_op.h"
#include "../mshadow_op.h"
#include "../elemwise_op_common.h"
#include "../tensor/init_op.h"
#include "./sampler.h"

namespace mxnet {
namespace op {


struct SampleOpParam {
  mxnet::TShape shape;
  std::string ctx;
  int dtype;
};

struct UniformParam {
  float low;
  float high;
};

struct NormalParam {
  float loc;
  float scale;
};

struct GammaParam {
  float alpha;
  float beta;
};

struct ExponentialParam {
  float lam;
};

struct PoissonParam {
  float lam;
};

struct NegBinomialParam {
  int k;
  float p;
};

struct GenNegBinomialParam {
  float mu;
  float alpha;
};

struct RandIntParam {
  int64_t low;
  int64_t high;
};

struct SampleUniformParam : public dmlc::Parameter<SampleUniformParam>,
  UniformParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleUniformParam) {
    DMLC_DECLARE_FIELD(low).set_default(0.0f)
    .describe("Lower bound of the distribution.");
    DMLC_DECLARE_FIELD(high).set_default(1.0f)
    .describe("Upper bound of the distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleNormalParam : public dmlc::Parameter<SampleNormalParam>,
  NormalParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleNormalParam) {
    DMLC_DECLARE_FIELD(loc).set_default(0.0f)
    .describe("Mean of the distribution.");
    DMLC_DECLARE_FIELD(scale).set_default(1.0f)
    .describe("Standard deviation of the distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleGammaParam : public dmlc::Parameter<SampleGammaParam>,
  GammaParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleGammaParam) {
    DMLC_DECLARE_FIELD(alpha).set_default(1.0f)
    .describe("Alpha parameter (shape) of the gamma distribution.");
    DMLC_DECLARE_FIELD(beta).set_default(1.0f)
    .describe("Beta parameter (scale) of the gamma distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleExponentialParam : public dmlc::Parameter<SampleExponentialParam>,
  ExponentialParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleExponentialParam) {
    DMLC_DECLARE_FIELD(lam).set_default(1.0f)
    .describe("Lambda parameter (rate) of the exponential distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SamplePoissonParam : public dmlc::Parameter<SamplePoissonParam>,
  PoissonParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SamplePoissonParam) {
    DMLC_DECLARE_FIELD(lam).set_default(1.0f)
    .describe("Lambda parameter (rate) of the Poisson distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleNegBinomialParam : public dmlc::Parameter<SampleNegBinomialParam>,
  NegBinomialParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleNegBinomialParam) {
    DMLC_DECLARE_FIELD(k).set_default(1)
    .describe("Limit of unsuccessful experiments.");
    DMLC_DECLARE_FIELD(p).set_default(1.0f)
    .describe("Failure probability in each experiment.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleGenNegBinomialParam : public dmlc::Parameter<SampleGenNegBinomialParam>,
  GenNegBinomialParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleGenNegBinomialParam) {
    DMLC_DECLARE_FIELD(mu).set_default(1.0f)
    .describe("Mean of the negative binomial distribution.");
    DMLC_DECLARE_FIELD(alpha).set_default(1.0f)
    .describe("Alpha (dispersion) parameter of the negative binomial distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("float32", mshadow::kFloat32)
    .add_enum("float64", mshadow::kFloat64)
    .add_enum("float16", mshadow::kFloat16)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleRandIntParam : public dmlc::Parameter<SampleRandIntParam>,
  RandIntParam, SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleRandIntParam) {
    DMLC_DECLARE_FIELD(low)
    .describe("Lower bound of the distribution.");
    DMLC_DECLARE_FIELD(high)
    .describe("Upper bound of the distribution.");
    DMLC_DECLARE_FIELD(shape)
    .set_default(mxnet::TShape())
    .describe("Shape of the output.");
    DMLC_DECLARE_FIELD(ctx)
    .set_default("")
    .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
              " Only used for imperative calls.");
    DMLC_DECLARE_FIELD(dtype)
    .add_enum("None", -1)
    .add_enum("int32", mshadow::kInt32)
    .add_enum("int64", mshadow::kInt64)
    .set_default(-1)
    .describe("DType of the output in case this can't be inferred. "
              "Defaults to int32 if not defined (dtype=None).");
  }
};

struct SampleUniformLikeParam : public dmlc::Parameter<SampleUniformLikeParam>,
  UniformParam {
  DMLC_DECLARE_PARAMETER(SampleUniformLikeParam) {
    DMLC_DECLARE_FIELD(low).set_default(0.0f)
    .describe("Lower bound of the distribution.");
    DMLC_DECLARE_FIELD(high).set_default(1.0f)
    .describe("Upper bound of the distribution.");
  }
};

struct SampleNormalLikeParam : public dmlc::Parameter<SampleNormalLikeParam>,
  NormalParam {
  DMLC_DECLARE_PARAMETER(SampleNormalLikeParam) {
    DMLC_DECLARE_FIELD(loc).set_default(0.0f)
    .describe("Mean of the distribution.");
    DMLC_DECLARE_FIELD(scale).set_default(1.0f)
    .describe("Standard deviation of the distribution.");
  }
};

struct SampleGammaLikeParam : public dmlc::Parameter<SampleGammaLikeParam>,
  GammaParam {
  DMLC_DECLARE_PARAMETER(SampleGammaLikeParam) {
    DMLC_DECLARE_FIELD(alpha).set_default(1.0f)
    .describe("Alpha parameter (shape) of the gamma distribution.");
    DMLC_DECLARE_FIELD(beta).set_default(1.0f)
    .describe("Beta parameter (scale) of the gamma distribution.");
  }
};

struct SampleExponentialLikeParam : public dmlc::Parameter<SampleExponentialLikeParam>,
  ExponentialParam {
  DMLC_DECLARE_PARAMETER(SampleExponentialLikeParam) {
    DMLC_DECLARE_FIELD(lam).set_default(1.0f)
    .describe("Lambda parameter (rate) of the exponential distribution.");
  }
};

struct SamplePoissonLikeParam : public dmlc::Parameter<SamplePoissonLikeParam>,
  PoissonParam {
  DMLC_DECLARE_PARAMETER(SamplePoissonLikeParam) {
    DMLC_DECLARE_FIELD(lam).set_default(1.0f)
    .describe("Lambda parameter (rate) of the Poisson distribution.");
  }
};

struct SampleNegBinomialLikeParam : public dmlc::Parameter<SampleNegBinomialLikeParam>,
  NegBinomialParam {
  DMLC_DECLARE_PARAMETER(SampleNegBinomialLikeParam) {
    DMLC_DECLARE_FIELD(k).set_default(1)
    .describe("Limit of unsuccessful experiments.");
    DMLC_DECLARE_FIELD(p).set_default(1.0f)
    .describe("Failure probability in each experiment.");
  }
};

struct SampleGenNegBinomialLikeParam : public dmlc::Parameter<SampleGenNegBinomialLikeParam>,
  GenNegBinomialParam {
  DMLC_DECLARE_PARAMETER(SampleGenNegBinomialLikeParam) {
    DMLC_DECLARE_FIELD(mu).set_default(1.0f)
    .describe("Mean of the negative binomial distribution.");
    DMLC_DECLARE_FIELD(alpha).set_default(1.0f)
    .describe("Alpha (dispersion) parameter of the negative binomial distribution.");
  }
};

using FSampleCompute = std::function<void (const nnvm::NodeAttrs& attrs,
                                           const OpContext& ctx,
                                           const OpReqType& req,
                                           TBlob* outputs)>;

using mxnet::TBlob;
using namespace mxnet::common::random;

// Allocates a single chunk of workspace memory and partitions it into three
// workspace tensors that hold the seeds as well as the distribution parameters.
template<typename xpu, typename DType>
MSHADOW_FORCE_INLINE void GetSamplingTempData(DType p1, DType p2, const OpContext& ctx,
                                              Tensor<xpu, 1, DType>* parm1,
                                              Tensor<xpu, 1, DType>* parm2) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  // Combined memory requirement for the workspace data.
  const index_t nInt((2 * sizeof(DType) + sizeof(unsigned) - 1) / sizeof(unsigned));
  Tensor<xpu, 1, unsigned> wspace
    = ctx.requested[1].get_space_typed<xpu, 1, unsigned>(Shape1(nInt), s);
  // Partition workspace into two chunks and initialize them.
  DType *pspace = static_cast<DType*>(static_cast<void*>(wspace.dptr_));
  *parm1 = Tensor<xpu, 1, DType>(pspace, Shape1(1), s);
  Copy(*parm1, Tensor<cpu, 1, DType>(&p1, Shape1(1)), s);
  *parm2 = Tensor<xpu, 1, DType>(pspace+1, Shape1(1), s);
  Copy(*parm2, Tensor<cpu, 1, DType>(&p2, Shape1(1)), s);
}

template<typename xpu, typename ParamType>
static inline void uniform_op(const nnvm::NodeAttrs& attrs,
                              const OpContext& ctx,
                              const OpReqType& req,
                              TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const UniformParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GE(param.high, param.low) << "low must be less or equal to high in uniform distribution";
  Tensor<xpu, 1, float> low, high;
  GetSamplingTempData<xpu, float>(param.low, param.high, ctx,
                                  &low, &high);
  UniformSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(low, high, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void normal_op(const nnvm::NodeAttrs& attrs,
                             const OpContext& ctx,
                             const OpReqType& req,
                             TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const NormalParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GT(param.scale, 0) << "scale parameter in gaussian has to be positive";
  Tensor<xpu, 1, float> loc, scale;
  GetSamplingTempData<xpu, float>(param.loc, param.scale, ctx, &loc, &scale);
  NormalSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(loc, scale, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void gamma_op(const nnvm::NodeAttrs& attrs,
                            const OpContext& ctx,
                            const OpReqType& req,
                            TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const GammaParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GT(param.alpha, 0) << "alpha parameter in gamma distribution has to be positive";
  CHECK_GT(param.beta, 0) << "beta parameter in gamma distribution has to be positive";
  Tensor<xpu, 1, float> alpha, beta;
  GetSamplingTempData<xpu, float>(param.alpha, param.beta, ctx, &alpha, &beta);
  GammaSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(alpha, beta, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void exponential_op(const nnvm::NodeAttrs& attrs,
                                  const OpContext& ctx,
                                  const OpReqType& req,
                                  TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const ExponentialParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GT(param.lam, 0) << "lambda parameter in exponential distribution has to be positive";
  Tensor<xpu, 1, float> lam, dummy;
  GetSamplingTempData<xpu, float>(param.lam, 0, ctx, &lam, &dummy);
  ExponentialSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(lam, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void poisson_op(const nnvm::NodeAttrs& attrs,
                              const OpContext& ctx,
                              const OpReqType& req,
                              TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const PoissonParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GE(param.lam, 0) << "lambda parameter in poisson distribution has to be non-negative";
  Tensor<xpu, 1, float> lam, dummy;
  GetSamplingTempData<xpu, float>(param.lam, 0, ctx, &lam, &dummy);
  PoissonSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(lam, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void neg_binomial_op(const nnvm::NodeAttrs& attrs,
                                   const OpContext& ctx,
                                   const OpReqType& req,
                                   TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const NegBinomialParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GE(param.k, 0) << "k parameter in negative binomial distribution has to be non-negative";
  CHECK_GE(param.p, 0) << "p parameter in negative binomial distribution has to be non-negative";
  Tensor<xpu, 1, float> k, p;
  GetSamplingTempData<xpu, float>(param.k, param.p, ctx, &k, &p);
  NegativeBinomialSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(k, p, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void gen_neg_binomial_op(const nnvm::NodeAttrs& attrs,
                                       const OpContext& ctx,
                                       const OpReqType& req,
                                       TBlob* outputs) {
  Stream<xpu> *s = ctx.get_stream<xpu>();
  const GenNegBinomialParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GE(param.mu, 0)
    << "mu parameter in generalized negative binomial distribution has to be non-negative";
  CHECK_GE(param.alpha, 0)
    << "alpha parameter in generalized negative binomial distribution has to be non-negative";
  Tensor<xpu, 1, float> mu, alpha;
  GetSamplingTempData<xpu, float>(param.mu, param.alpha, ctx, &mu, &alpha);
  GeneralizedNegativeBinomialSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH(outputs[0].type_flag_, OType, {
    RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
    Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
    sampler.Sample(mu, alpha, out, pgen, s);
  });
}

template<typename xpu, typename ParamType>
static inline void rand_int_op(const nnvm::NodeAttrs& attrs,
                               const OpContext& ctx,
                               const OpReqType& req,
                               TBlob* outputs) {
    Stream<xpu> *s = ctx.get_stream<xpu>();
    const SampleRandIntParam& param = nnvm::get<SampleRandIntParam>(attrs.parsed);
    CHECK_GE(param.high, param.low) << "low must be less or equal to high in uniform distribution";
    Tensor<xpu, 1, int64_t> low, high;
    GetSamplingTempData<xpu, int64_t>(param.low, param.high, ctx,
                                    &low, &high);
    RandIntSampler<xpu> sampler;
    MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, OType, {
      RandGenerator<xpu, OType> *pgen = ctx.requested[0].get_parallel_random<xpu, OType>();
      Tensor<xpu, 1, OType> out = outputs->FlatTo1D<xpu, OType>(s);
      sampler.Sample(low, high, out, pgen, s);
    });
}

template<typename xpu, typename ParamType>
struct SampleMaster;

template<typename xpu>
struct SampleMaster<xpu, SampleUniformParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    uniform_op<xpu, SampleUniformParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleUniformLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    uniform_op<xpu, SampleUniformLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleNormalParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    normal_op<xpu, SampleNormalParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleNormalLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    normal_op<xpu, SampleNormalLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleGammaParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    gamma_op<xpu, SampleGammaParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleGammaLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    gamma_op<xpu, SampleGammaLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleExponentialParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    exponential_op<xpu, SampleExponentialParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleExponentialLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    exponential_op<xpu, SampleExponentialLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SamplePoissonParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    poisson_op<xpu, SamplePoissonParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SamplePoissonLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    poisson_op<xpu, SamplePoissonLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleNegBinomialParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    neg_binomial_op<xpu, SampleNegBinomialParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleNegBinomialLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    neg_binomial_op<xpu, SampleNegBinomialLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleGenNegBinomialParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    gen_neg_binomial_op<xpu, SampleGenNegBinomialParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleGenNegBinomialLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    gen_neg_binomial_op<xpu, SampleGenNegBinomialLikeParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu>
struct SampleMaster<xpu, SampleRandIntParam> {
  static void op(const nnvm::NodeAttrs& attrs,
                 const OpContext& ctx,
                 const OpReqType& req,
                 TBlob* outputs) {
    rand_int_op<xpu, SampleRandIntParam>(attrs, ctx, req, outputs);
  }
};

template<typename xpu, typename ParamType>
void SampleComputeEx_(const nnvm::NodeAttrs& attrs,
                      const OpContext& ctx,
                      const std::vector<NDArray>& inputs,
                      const std::vector<OpReqType>& req,
                      const std::vector<NDArray>& outputs,
                      SampleMaster<xpu, ParamType> sample_master) {
  using namespace mxnet::op;
  NDArray output = outputs[0];
  mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
  if (output.storage_type() == kRowSparseStorage) {
    // indices
    nnvm::dim_t nnr = output.shape()[0];
    output.CheckAndAlloc({mshadow::Shape1(nnr)});
    MSHADOW_IDX_TYPE_SWITCH(output.aux_type(rowsparse::kIdx), IType, {
      IType* idx = output.aux_data(rowsparse::kIdx).dptr<IType>();
      mxnet_op::Kernel<PopulateFullIdxRspKernel, xpu>::Launch(s, nnr, idx);
    });
    // data
    TBlob out_blob = output.data();
    sample_master.op(attrs, ctx, req[0], &out_blob);
  } else {
    LOG(FATAL) << "Unexpected storage type for SampleComputeEx_: "
               << output.storage_type();
  }
}

template<typename xpu, typename ParamType>
void Sample_(const nnvm::NodeAttrs& attrs,
             const OpContext& ctx,
             const std::vector<TBlob>& inputs,
             const std::vector<OpReqType>& req,
             const std::vector<TBlob>& outputs) {
  TBlob out = outputs[0];
  SampleMaster<xpu, ParamType>::op(attrs, ctx, req[0], &out);
}

template<typename xpu, typename ParamType>
void SampleEx_(const nnvm::NodeAttrs& attrs,
               const OpContext& ctx,
               const std::vector<NDArray>& inputs,
               const std::vector<OpReqType>& req,
               const std::vector<NDArray>& outputs) {
  SampleMaster<xpu, ParamType> sample_master;
  SampleComputeEx_<xpu, ParamType>(attrs, ctx, inputs, req, outputs, sample_master);
}

template<typename ParamType>
inline bool SampleOpType(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(), 0);
  CHECK_EQ(out_type->size(), 1);
  int dtype = -1;
  int dtype_out = (*out_type)[0];
  if (dtype_out != -1) {
    // Output type can be inferred, use it and make sure it matches
    dtype = dtype_out;
    if (param.dtype != -1) {
      // dtype given in args, check that it matches the output type
      CHECK_EQ(dtype_out, param.dtype) << "Output type does not match requested type: "
      << dtype_out << " vs " << param.dtype;
    }
  } else {
    // Output type can't be inferred
    if (param.dtype != -1) {
      // Use dtype given in args
      dtype = param.dtype;
    } else {
      // Use default
      dtype = mxnet::common::GetDefaultDtype();
    }
  }
  bool dtype_ok = (dtype == mshadow::kFloat16) || (dtype == mshadow::kFloat32) ||
  (dtype == mshadow::kFloat64);
  CHECK(dtype_ok) << "Output type must be float16, float32, float64: dtype is "
  << dtype_out << " vs " << mshadow::kFloat16 << " or " << mshadow::kFloat32 << " or "
  << mshadow::kFloat64;
  TYPE_ASSIGN_CHECK(*out_type, 0, dtype);
  return true;
}

template<>
inline bool SampleOpType<SampleRandIntParam>(const nnvm::NodeAttrs& attrs,
                         std::vector<int> *in_type,
                         std::vector<int> *out_type) {
  const SampleRandIntParam& param = nnvm::get<SampleRandIntParam>(attrs.parsed);
  CHECK_EQ(in_type->size(), 0);
  CHECK_EQ(out_type->size(), 1);
  int dtype = -1;
  int dtype_out = (*out_type)[0];
  if (dtype_out != -1) {
    // Output type can be inferred, use it and make sure it matches
    dtype = dtype_out;
    if (param.dtype != -1) {
      // dtype given in args, check that it matches the output type
      CHECK_EQ(dtype_out, param.dtype) << "Output type does not match requested type: "
      << dtype_out << " vs " << param.dtype;
    }
  } else {
    // Output type can't be inferred
    if (param.dtype != -1) {
      // Use dtype given in args
      dtype = param.dtype;
    } else {
      // Use default
      dtype = mshadow::kInt32;
    }
  }
  bool dtype_ok = (dtype == mshadow::kInt32) || (dtype == mshadow::kInt64);
  CHECK(dtype_ok) << "Output type must be int32, int64: dtype is "
  << dtype_out << " vs " << mshadow::kInt32 << " or " << mshadow::kInt64;
  TYPE_ASSIGN_CHECK(*out_type, 0, dtype);
  return true;
}

inline std::vector<ResourceRequest> SampleResource(const NodeAttrs& attrs) {
  return { ResourceRequest::kParallelRandom, ResourceRequest::kTempSpace };
}

}  // namespace op
}  // namespace mxnet
#endif  // MXNET_OPERATOR_RANDOM_SAMPLE_OP_H_