title: using theano randomstream choice for a batch of 1D distributions (2D matrix)?
body: <p>I have a matrix in which each row is a distribution. I want to sample an index from each row of this matrix in theano.</p>
  
  <p>I am able to sample from a single vector, see function <code>f</code> in the code below.
  But I am not able to extend this to a matrix, <code>f_mat</code> does not work.</p>
  
  <pre><code>import theano
  import theano.tensor as tt
  import theano.tensor.shared_randomstreams
  import numpy as np
  
  n = 3
  seed = 1
  
  rng = tt.shared_randomstreams.RandomStreams(seed=seed)
  
  dist1 = tt.fvector('dist')
  i = rng.choice(size=(1,), a=dist1.shape[0], p=dist1)
  f = theano.function([dist1], [i])
  _dist = np.asarray([0.5, 0.5, 0.])
  print f(_dist.astype(np.float32)) #THIS IS WORKING
  
  dist_mat = tt.fmatrix('dist_mat')
  i_mat = rng.choice(size=(dist_mat.shape[0],), a=dist_mat.shape[1], p=dist_mat, ndim=1)
  f_mat = theano.function([dist_mat], [i_mat])
  _dist_mat = np.asarray([[0.5, 0.1, 0.4], [0.1, 0.1, 0.8]])
  print f_mat(_dist_mat.astype(np.float32)) #THIS DOES NOT WORK
  </code></pre>
  
  <p>What am I doing wrong? I would like <code>f_mat</code> to produce a vector with the same dimensionality as the number of rows in the <code>_dist_mat</code> matrix.</p>