using theano randomstream choice for a batch of 1D distributions (2D matrix)?

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.

I am able to sample from a single vector, see function f in the code below. But I am not able to extend this to a matrix, f_mat does not work.

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

What am I doing wrong? I would like f_mat to produce a vector with the same dimensionality as the number of rows in the _dist_mat matrix.