text: # TODO test dtype != float32
  from __future__ import absolute_import, print_function, division
  import os
  import warnings
  
  try:
      import pygpu
  except ImportError:
      pass
  
  import theano
  import theano.sandbox.multinomial
  from theano import Apply, config
  from theano.gof import Op
  
  from theano.tensor import NotScalarConstantError, get_scalar_constant_value
  from .basic_ops import as_gpuarray_variable, infer_context_name, GpuKernelBase, Kernel
  from .opt import register_opt, op_lifter, register_opt2
  from .type import GpuArrayType
  from .elemwise import GpuDimShuffle
  from theano.scalar import as_scalar
  from .fp16_help import write_w, load_w, work_dtype
  
  
  class GPUAMultinomialFromUniform(GpuKernelBase, Op):
      __props__ = ("odtype",)
      _f16_ok = True
  
      def __init__(self, odtype):
          Op.__init__(self)
          self.odtype = odtype
  
      def get_params(self, node):
          return node.outputs[0].type.context
  
      def c_headers(self):
          return ['<numpy_compat.h>', 'gpuarray_helper.h']
  
      def c_header_dirs(self):
          return [os.path.dirname(__file__)]
  
      def make_node(self, pvals, unis):
          assert unis.dtype == pvals.dtype
          assert pvals.dtype in ['float32', 'float16', 'float64']
          ctx_name = infer_context_name(pvals, unis)
          pvals = as_gpuarray_variable(pvals, ctx_name)
          unis = as_gpuarray_variable(unis, ctx_name)
  
          if pvals.ndim != 2:
              raise NotImplementedError('pvals ndim should be 2', pvals.ndim)
          if unis.ndim != 1:
              raise NotImplementedError('unis ndim should be 1', unis.ndim)
          if self.odtype == 'auto':
              odtype = pvals.dtype
          else:
              odtype = self.odtype
          br = (pvals.broadcastable[1], pvals.broadcastable[0])
          out = GpuArrayType(broadcastable=br,
                             dtype=odtype,
                             context_name=ctx_name)()
  
          return Apply(self, [pvals, unis], [out])
  
      def gpu_kernels(self, node, name):
          out_ctype = pygpu.gpuarray.dtype_to_ctype(node.outputs[0].dtype)
          in_ctype = pygpu.gpuarray.dtype_to_ctype(node.inputs[0].dtype)
          work_ctype = pygpu.gpuarray.dtype_to_ctype(work_dtype(node.inputs[0].dtype))
          write_out_ctype = write_w(node.outputs[0].dtype)
          load_in_ctype = load_w(node.inputs[0].dtype)
          code = """
  KERNEL void k_multi_warp_multinomial(
      const ga_size nb_multi,
      const ga_size nb_outcomes,
      GLOBAL_MEM %(in_ctype)s * global_pvals,
      const ga_ssize pvals_row_stride,
      const ga_ssize pvals_col_stride,
      GLOBAL_MEM %(in_ctype)s * global_unis,
      const ga_ssize unis_stride,
      GLOBAL_MEM %(out_ctype)s * global_outs,
      const ga_ssize outs_row_stride,
      const ga_ssize outs_col_stride
  )
  {
      // each thread takes care of one multinomial draw
      int n = LDIM_0*GID_0 + LID_0;
      if (n < nb_multi)
      {
          %(work_ctype)s cummul = 0.;
          bool done = false;
          const %(work_ctype)s unis_n = %(load_in_ctype)s(global_unis[n*unis_stride]);
          for (ga_size m = 0; m < nb_outcomes; ++m)
          {
              %(work_ctype)s current_out = 0;
              if (!done)
              {
                  cummul += %(load_in_ctype)s(global_pvals[m * pvals_col_stride + n * pvals_row_stride]);
                  if (unis_n < cummul)
                  {
                      current_out = 1;
                      done = true;
                  }
              }
              //write out transposed for speed.
              global_outs[n * outs_col_stride +
                          m * outs_row_stride] = %(write_out_ctype)s(current_out);
          }
      }
  }
  """ % dict(out_ctype=out_ctype, write_out_ctype=write_out_ctype,
             work_ctype=work_ctype, in_ctype=in_ctype, load_in_ctype=load_in_ctype)
          return [Kernel(
              code=code, name="k_multi_warp_multinomial",
              params=[pygpu.gpuarray.SIZE,
                      pygpu.gpuarray.SIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.SSIZE],
              flags=Kernel.get_flags(node.outputs[0].dtype),
              objvar='k_multi_warp_multinomial_' + name)]
  
      def c_code(self, node, name, inp, outputs, sub):
          pvals, unis = inp
          out, = outputs
          fail = sub['fail']
          ctx = sub['params']
          sync = bool(config.gpuarray.sync)
          kname = self.gpu_kernels(node, name)[0].objvar
          out_typecode = pygpu.gpuarray.dtype_to_typecode(node.outputs[0].dtype)
          in_typecode = pygpu.gpuarray.dtype_to_typecode(node.inputs[0].dtype)
          s = """
          PyGpuArrayObject * pvals = %(pvals)s;
          PyGpuArrayObject * unis = %(unis)s;
          PyGpuArrayObject * out = %(out)s;
  
      size_t dims[2];
      if (PyGpuArray_NDIM(pvals) != 2)
      {
          PyErr_Format(PyExc_TypeError, "pvals wrong rank");
          %(fail)s
      }
      if (PyGpuArray_NDIM(unis) != 1)
      {
          PyErr_Format(PyExc_TypeError, "unis wrong rank");
          %(fail)s
      }
      if (PyGpuArray_DIMS(unis)[0] != PyGpuArray_DIMS(pvals)[0])
      {
          PyErr_Format(PyExc_ValueError, "unis.shape[0] != pvals.shape[0]");
          %(fail)s
      }
  
      dims[0] = PyGpuArray_DIMS(pvals)[1];
      dims[1] = PyGpuArray_DIMS(pvals)[0];
      if (theano_prep_output(&out, 2, dims, %(out_typecode)s,
                             GA_C_ORDER, %(ctx)s) != 0){
        %(fail)s
      }
      %(out)s = out;
      GpuArray_memset(&(out->ga), 0);
      { // NESTED SCOPE
          int nb_multi = PyGpuArray_DIMS(pvals)[0];
          int nb_outcomes = PyGpuArray_DIMS(pvals)[1];
          //TODO : change this for a beautiful constant
          int max_nb_blocks = 2<<15 - 1;
          size_t nb_blocks = max_nb_blocks + 1;
          size_t nb_threads=16; // so it really starts at 32, because of the *2
          do
          {
              nb_threads*=2;
              if (nb_multi % %nb_threads == 0)
                  nb_blocks = nb_multi/nb_threads;
              else
                  nb_blocks = (int)((float)nb_multi/(float)nb_threads + 1.);
          } while (nb_blocks > max_nb_blocks);
  
          //printf("\\nN=%%i b=%%i t=%%i t*b=%%i",
          //         nb_multi, nb_blocks, nb_threads, nb_blocks*nb_threads);
  
          // TODO : next line is a bit hardcoded...
          if (nb_threads > 512)
          {
              PyErr_Format(
                  PyExc_ValueError,
                  "Multinomial is not implemented for so many rows in the matrix (%%i)",
                  nb_multi);
              %(fail)s
          }
  
          assert(nb_blocks*nb_threads >= nb_multi);
  
          void *args[10];
          ssize_t strides[5] = {
              PyGpuArray_STRIDES(pvals)[0]/gpuarray_get_elsize(%(in_typecode)s),
              PyGpuArray_STRIDES(pvals)[1]/gpuarray_get_elsize(%(in_typecode)s),
              PyGpuArray_STRIDES(unis)[0]/gpuarray_get_elsize(%(in_typecode)s),
              PyGpuArray_STRIDES(out)[0]/gpuarray_get_elsize(%(out_typecode)s),
              PyGpuArray_STRIDES(out)[1]/gpuarray_get_elsize(%(out_typecode)s)
          };
          int err;
          args[0] = (void*)&PyGpuArray_DIMS(out)[1];
          args[1] = (void*)&PyGpuArray_DIMS(out)[0];
          args[2] = pvals->ga.data; //PyGpuArray_DEV_DATA(pvals);
          args[3] = (void*)&strides[0];
          args[4] = (void*)&strides[1];
          args[5] = unis->ga.data; //PyGpuArray_DEV_DATA(unis);
          args[6] = (void*)&strides[2];
          args[7] = out->ga.data; //PyGpuArray_DEV_DATA(out);
          args[8] = (void*)&strides[3];
          args[9] = (void*)&strides[4];
  
          err = GpuKernel_call(&%(kname)s, 1, &nb_blocks, &nb_threads, 0, args);
          if (err != GA_NO_ERROR) {
             PyErr_Format(
                  PyExc_RuntimeError,
                  "gpuarray error: %%s: %%s.\\n",
                  "k_multi_warp_%(name)s",
                  GpuKernel_error(&%(kname)s, err));
              %(fail)s;
          }
          if(%(sync)d)
              GpuArray_sync(&(out->ga));
      } // END NESTED SCOPE
          """ % locals()
  
          return s
  
      def c_code_cache_version(self):
          return (3,)
  
  
  class GPUAChoiceFromUniform(GpuKernelBase, Op):
      """
      The output is transposed compared to MultinomialWOReplacementFromUniform.
      We must insert a Transpose op after it.
  
      The optimization that moves it to the gpu does it.
  
      """
  
      __props__ = ("odtype", "replace")
  
      def __init__(self, odtype, replace=False):
          Op.__init__(self)
          self.odtype = odtype
          self.replace = replace
  
      def __setstate__(self, state):
          self.__dict__.update(state)
          if "replace" not in state:
              self.replace = False
  
      def get_params(self, node):
          return node.outputs[0].type.context
  
      def c_headers(self):
          return ['<numpy_compat.h>', 'gpuarray_helper.h']
  
      def c_header_dirs(self):
          return [os.path.dirname(__file__)]
  
      def make_node(self, pvals, unis, n):
          assert pvals.dtype == 'float32'
          assert unis.dtype == 'float32'
          ctx_name = infer_context_name(pvals, unis)
  
          pvals = as_gpuarray_variable(pvals, ctx_name)
          unis = as_gpuarray_variable(unis, ctx_name)
  
          if pvals.ndim != 2:
              raise NotImplementedError('pvals ndim should be 2', pvals.ndim)
          if unis.ndim != 1:
              raise NotImplementedError('unis ndim should be 1', unis.ndim)
          if self.odtype == 'auto':
              odtype = 'int64'
          else:
              odtype = self.odtype
          assert odtype == 'int64', odtype
          br = (pvals.broadcastable[1], pvals.broadcastable[0])
          out = GpuArrayType(broadcastable=br,
                             dtype=odtype,
                             context_name=ctx_name)()
  
          return Apply(self, [pvals, unis, as_scalar(n)], [out])
  
      def gpu_kernels(self, node, name):
          replace = int(self.replace)
          code = """
  KERNEL void k_multi_warp_multinomial_wor(
      const ga_size nb_multi,
      const ga_size nb_outcomes,
      const ga_size n_samples,
      GLOBAL_MEM float * global_pvals_copy,
      const ga_ssize pvals_row_stride,
      const ga_ssize pvals_col_stride,
      GLOBAL_MEM float * global_unis,
      const ga_ssize unis_stride,
      GLOBAL_MEM ga_long * global_outs,
      const ga_ssize outs_row_stride,
      const ga_ssize outs_col_stride
  )
  {
      // each thread takes care of one multinomial-wor n_samples-draw
      int n = LDIM_0*GID_0 + LID_0;
  
      if (n < nb_multi)
      {
          for (int c = 0; c < n_samples; ++c)
          {
              float cummul = 0.;
              bool done = false;
              const float unis_n = global_unis[(c * nb_multi + n)*unis_stride];
              for (ga_size m = 0; m < nb_outcomes; ++m)
              {
                  float pvals_nm = global_pvals_copy[m * pvals_col_stride + n * pvals_row_stride];
                  cummul += pvals_nm;
  
                  if (!done && unis_n < cummul)
                  {
                      //write out transposed for speed.
                      global_outs[n * outs_col_stride +
                                  c * outs_row_stride] = m;
  
                      if (! %(replace)s )
                      {
                          global_pvals_copy[m * pvals_col_stride + n * pvals_row_stride] = 0.0;
                          cummul -= pvals_nm;
                      }
                      done = true;
                  }
              }
              // No need to renormalize after the last samples.
              if (c == (n_samples - 1))
                  break;
              if (! %(replace)s )
              {
                  // parallel renormalize the multinomial
                  for (ga_int k = LID_1; k < nb_outcomes; k+=LDIM_1)
                  {
                      global_pvals_copy[k * pvals_col_stride + n * pvals_row_stride] /= cummul;
                  }
              }
          }
      }
  }
  """ % {"replace": replace}
          return [Kernel(
              code=code, name="k_multi_warp_multinomial_wor",
              params=[pygpu.gpuarray.SIZE,
                      pygpu.gpuarray.SIZE,
                      pygpu.gpuarray.SIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.GpuArray,
                      pygpu.gpuarray.SSIZE,
                      pygpu.gpuarray.SSIZE
                      ],
              flags=Kernel.get_flags(node.outputs[0].dtype),
              objvar='k_multi_warp_multinomial_wor_' + name)]
  
      def c_code(self, node, name, inp, outputs, sub):
          pvals, unis, n = inp
          out, = outputs
          replace = int(self.replace)
          fail = sub['fail']
          ctx = sub['params']
          sync = bool(config.gpuarray.sync)
          kname = self.gpu_kernels(node, name)[0].objvar
          s = """
      PyGpuArrayObject * pvals = %(pvals)s;
      PyGpuArrayObject * unis = %(unis)s;
      const size_t n_samples = %(n)s;
      PyGpuArrayObject * out = %(out)s;
      // create a copy of pvals matrix
      PyGpuArrayObject * pvals_copy = NULL;
  
      size_t dims[2];
      if (PyGpuArray_NDIM(pvals) != 2)
      {
          PyErr_Format(PyExc_TypeError, "pvals wrong rank");
          %(fail)s
      }
      if (PyGpuArray_NDIM(unis) != 1)
      {
          PyErr_Format(PyExc_TypeError, "unis wrong rank");
          %(fail)s
      }
      if ( n_samples > (PyGpuArray_DIMS(pvals)[1]) )
      {
          PyErr_Format(PyExc_ValueError, "Cannot sample without replacement n samples bigger than the size of the distribution.");
          %(fail)s;
      }
      if (PyGpuArray_DIMS(unis)[0] != PyGpuArray_DIMS(pvals)[0] * n_samples)
      {
          PyErr_Format(PyExc_ValueError, "unis.shape[0] != pvals.shape[0] * n");
          %(fail)s
      }
  
      pvals_copy = pygpu_copy(pvals, GA_C_ORDER);
  
      dims[0] = n_samples;
      dims[1] = PyGpuArray_DIMS(pvals)[0];
  
      if (theano_prep_output(&out, 2, dims, GA_LONG,
                             GA_C_ORDER, %(ctx)s) != 0){
          Py_DECREF(pvals_copy);
          %(fail)s
      }
  
      %(out)s = out;
  
      { // NESTED SCOPE
          int nb_multi = PyGpuArray_DIMS(pvals)[0];
          int nb_outcomes = PyGpuArray_DIMS(pvals)[1];
          //TODO : change this for a beautiful constant
          int max_nb_blocks = 2<<15 - 1;
          size_t nb_blocks = max_nb_blocks + 1;
          size_t nb_threads=16; // so it really starts at 32, because of the *2
          do
          {
              nb_threads*=2;
              if (nb_multi %% nb_threads == 0)
                  nb_blocks = nb_multi/nb_threads;
              else
                  nb_blocks = (int)((float)nb_multi/(float)nb_threads + 1.);
          } while (nb_blocks > max_nb_blocks);
  
          // TODO : next line is a bit hardcoded...
          if (nb_threads > 512)
          {
              PyErr_Format(
                  PyExc_ValueError,
                  "Multinomial is not implemented for so many rows in the matrix (%%i)",
                  nb_multi);
              Py_DECREF(pvals_copy);
              %(fail)s
          }
  
          assert(nb_blocks*nb_threads >= nb_multi);
  
          void *args[11];
          ssize_t strides[5] = {
              PyGpuArray_STRIDES(pvals)[0]/sizeof(float),
              PyGpuArray_STRIDES(pvals)[1]/sizeof(float),
              PyGpuArray_STRIDES(unis)[0]/sizeof(float),
              PyGpuArray_STRIDES(out)[0]/8,
              PyGpuArray_STRIDES(out)[1]/8
          };
          int err;
          args[0] = (void*)&PyGpuArray_DIMS(pvals)[0];
          args[1] = (void*)&PyGpuArray_DIMS(pvals)[1];
          args[2] = (void*)&n_samples;
          args[3] = pvals_copy->ga.data; //PyGpuArray_DEV_DATA(pvals);
          args[4] = (void*)&strides[0];
          args[5] = (void*)&strides[1];
          args[6] = unis->ga.data; //PyGpuArray_DEV_DATA(unis);
          args[7] = (void*)&strides[2];
          args[8] = out->ga.data; //PyGpuArray_DEV_DATA(out);
          args[9] = (void*)&strides[3];
          args[10] = (void*)&strides[4];
  
          size_t nb_threads2[2], nb_blocks2[2];
          nb_threads2[0] = nb_threads;
          nb_threads2[1] = 1;
          // If we can't schedule enough threads parallelize the renormalization.
          // I do this because we don't always use those extra threads.
          if ((nb_threads * nb_blocks < 2048) && ! %(replace)d )
              nb_threads2[1] = 1024 / nb_threads;
  
          nb_blocks2[0] = nb_blocks;
          nb_blocks2[1] = 1;
  
          err = GpuKernel_call(&%(kname)s, 2, nb_blocks2, nb_threads2, 0, args);
          if (err != GA_NO_ERROR) {
             PyErr_Format(
                  PyExc_RuntimeError,
                  "gpuarray error: %%s: %%s.\\n",
                  "k_multi_warp_%(name)s",
                  GpuKernel_error(&%(kname)s, err));
             Py_DECREF(pvals_copy);
             %(fail)s;
          }
          if(%(sync)d)
              GpuArray_sync(&(out->ga));
          Py_DECREF(pvals_copy);
      } // END NESTED SCOPE
          """ % locals()
          return s
  
      def c_code_cache_version(self):
          return (4,)
  
  
  @register_opt('fast_compile')
  @op_lifter([theano.sandbox.multinomial.MultinomialFromUniform])
  @register_opt2([theano.sandbox.multinomial.MultinomialFromUniform], 'fast_compile')
  def local_gpua_multinomial(op, context_name, inputs, outputs):
      # TODO : need description for function
  
      if len(inputs) == 2:
          p, u = inputs
          n_samples = 1
      else:
          p, u, n_samples = inputs
      try:
          if get_scalar_constant_value(n_samples) != 1:
              return None
      except NotScalarConstantError:
          return None
      m, = outputs
      gpu_op = GPUAMultinomialFromUniform(op.odtype)
      return GpuDimShuffle([False, False], [1, 0])(
          gpu_op(p, u))
  
  
  @register_opt('fast_compile')
  @op_lifter([theano.sandbox.multinomial.ChoiceFromUniform])
  @register_opt2([theano.sandbox.multinomial.ChoiceFromUniform], 'fast_compile')
  def local_gpua_multinomial_wor(op, context_name, inputs, outputs):
      # TODO : need description for function
      p, u, n = inputs
      m, = outputs
      if ((p.dtype == u.dtype == 'float32') and (m.dtype == 'int64')):
          gpu_op = GPUAChoiceFromUniform(op.odtype)
          return GpuDimShuffle([False, False], [1, 0])(
              gpu_op(p, u, n))
  
  
  class GPUAMultinomialWOReplacementFromUniform(GPUAChoiceFromUniform):
      def __init__(self, *args, **kwargs):
          warnings.warn("GPUAMultinomialWOReplacementFromUniform is deprecated, "
                        "use GPUAChoiceFromUniform instead.",
                        DeprecationWarning,
                        stacklevel=2)
          super(GPUAMultinomialWOReplacementFromUniform, self).__init__(*args, **kwargs)
  