numpy - argmax result as a subtensor -

i want use argmax kept dimensions subtensor. have:

m, argm = t.max_and_argmax(a, axis=axis, keepdims=true) 

and want set values 0 in a. i.e. need use t.set_subtensor. use that, need specify subtensor a_sub of a @ argm i'm not sure how looks like. a_sub = a[argm] wrong multiple dimensions.

this should hold:

a_sub == t.max(a, axis=axis) a_sub.shape == t.max(a, axis=axis).shape 

in end, want do:

a = t.set_subtensor(a_sub, 0) 

my current solution:

idx = t.arange(a.shape[axis]).dimshuffle(['x'] * axis + [0] + ['x'] * (a.ndim - axis - 1)) = t.switch(t.eq(idx, argm), 0, a) 

however, a_sub = a[t.eq(idx, argm)] not work.

you need use theano's advanced indexing features which, unfortunately, differ numpy's advanced indexing.

here's example want.

update: works parametrized axis note axis cannot symbolic.

import numpy  import theano import theano.tensor tt  theano.config.compute_test_value = 'raise'  axis = 2  x = tt.tensor3() x.tag.test_value = numpy.array([[[3, 2, 6], [5, 1, 4]], [[2, 1, 6], [6, 1, 5]]],                                dtype=theano.config.floatx)  # identify largest value in each row x_argmax = tt.argmax(x, axis=axis, keepdims=true)  # construct row of indexes length of axis indexes = tt.arange(x.shape[axis]).dimshuffle(     *(['x' dim1 in xrange(axis)] + [0] + ['x' dim2 in xrange(x.ndim - axis - 1)]))  # create binary mask indicating maximum values appear mask = tt.eq(indexes, x_argmax)  # alter original matrix @ places maximum values appeared x_prime = tt.set_subtensor(x[mask.nonzero()], 0)  print x_prime.tag.test_value