src/operator/random/sample_op.h - master - githubmirror/incubator-mxnet

src/operator/random/sample_op.h
859 строк34 KB
Paweł Głomski
Improve bf16 support (#21002)
15 июл 2022, 11:56
15 июл 2022, 11:56f6d1ed1
/*
 * 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 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 BinomialParam {
  int n;
  float p;
};

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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .set_default(-1)
        .describe(
            "DType of the output in case this can't be inferred. "
            "Defaults to float32 if not defined (dtype=None).");
  }
};

struct SampleBinomialParam : public dmlc::Parameter<SampleBinomialParam>,
                             BinomialParam,
                             SampleOpParam {
  DMLC_DECLARE_PARAMETER(SampleBinomialParam) {
    DMLC_DECLARE_FIELD(n).set_default(1).describe("number of experiments.");
    DMLC_DECLARE_FIELD(p).set_default(1.0f).describe("success 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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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)
        .add_enum("bfloat16", mshadow::kBfloat16)
        .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).");
  }
  void SetAttrDict(std::unordered_map<std::string, std::string>* dict) {
    std::ostringstream low_s, high_s, dtype_s, shape_s;
    low_s << low;
    high_s << high;
    dtype_s << dtype;
    shape_s << shape;
    (*dict)["low"]   = low_s.str();
    (*dict)["high"]  = high_s.str();
    (*dict)["dtype"] = MXNetTypeWithBool2String(dtype);
    (*dict)["shape"] = shape_s.str();
  }
};

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 SampleBinomialLikeParam : public dmlc::Parameter<SampleBinomialLikeParam>, BinomialParam {
  DMLC_DECLARE_PARAMETER(SampleBinomialLikeParam) {
    DMLC_DECLARE_FIELD(n).set_default(1).describe("Number of experiments.");
    DMLC_DECLARE_FIELD(p).set_default(1.0f).describe("success probability in each experiment.");
  }
};

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_EX(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_EX(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_EX(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_EX(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_EX(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 binomial_op(const nnvm::NodeAttrs& attrs,
                               const OpContext& ctx,
                               const OpReqType& req,
                               TBlob* outputs) {
  Stream<xpu>* s             = ctx.get_stream<xpu>();
  const BinomialParam& param = nnvm::get<ParamType>(attrs.parsed);
  CHECK_GE(param.n, 0) << "n parameter in binomial distribution has to be non-negative";
  CHECK_GE(param.p, 0) << "p parameter in binomial distribution has to be non-negative";
  Tensor<xpu, 1, float> n, p;
  GetSamplingTempData<xpu, float>(param.n, param.p, ctx, &n, &p);
  BinomialSampler<xpu> sampler;
  MSHADOW_REAL_TYPE_SWITCH_EX(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(n, p, 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_EX(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_EX(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, SampleBinomialParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    binomial_op<xpu, SampleBinomialParam>(attrs, ctx, req, outputs);
  }
};

template <typename xpu>
struct SampleMaster<xpu, SampleBinomialLikeParam> {
  static inline void op(const nnvm::NodeAttrs& attrs,
                        const OpContext& ctx,
                        const OpReqType& req,
                        TBlob* outputs) {
    binomial_op<xpu, SampleBinomialLikeParam>(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::kBfloat16 || dtype == mshadow::kFloat16 ||
                  dtype == mshadow::kFloat32 || dtype == mshadow::kFloat64;
  CHECK(dtype_ok) << "Output type must be bfloat16, float16, float32, float64: dtype is "
                  << dtype_out << " vs " << mshadow::kBfloat16 << " or " << 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_