I have a problem with Theano, namely the Jacobian of a Theano function found by auto-differentiation, returns all NaNs. I suspect the problem might be the switch statements because the gradient calculation worked for a similar function that didn't have any switch statements, but in theory the switch function should not break the gradient. I was wondering what could be causing this.
The following function returns the Theano function from a bunch of inputs to the Jacobian of the flattened s3 variable: ```python def theano_RSA( possible_signals_array=T.ltensor3("possible_signals_array"), real_signals_indices=T.lvector("real_signals_indices"), alphas=T.dvector("alphas"), choice_alphas=T.dvector("choice_alphas"), cost_factors= T.dvector("cost_factors"), objective_costs_possible=T.dvector("objective_costs_possible"), at_most_as_costly=T.lmatrix("at_most_as_costly"), types=T.lvector("types"), distances=T.dmatrix("distances")): real_signals_array = possible_signals_array[:,real_signals_indices] objective_costs_real = objective_costs_possible[real_signals_indices] max_pic_size = possible_signals_array.shape[0] - 1 considered_signals = at_most_as_costly & types.dimshuffle("x", 0) language_l = possible_signals_array / possible_signals_array.sum(axis=-1, keepdims=True) expected_dist_l0 = T.tensordot(language_l, distances, axes=[[2],[0]]) unnorm_l0 = T.exp(choice_alphas[:,np.newaxis,np.newaxis,np.newaxis]*-expected_dist_l0) shape = unnorm_l0.shape _, picsize_index, _, state_index = T.mgrid[ 0:shape[0], 0:shape[1], 0:shape[2], 0:shape[3] ] unnorm_l0 = T.switch(state_index > picsize_index, 0, unnorm_l0) l0 = unnorm_l0 / T.sum(unnorm_l0, axis=-1, keepdims=True) l0_extended = l0[:,:,np.newaxis,:,:] costs_possible = T.outer(cost_factors, objective_costs_possible) utility_l0 = T.log(l0_extended) unnorm_s1 = T.exp( alphas[:,np.newaxis, np.newaxis, np.newaxis, np.newaxis] * (utility_l0 - costs_possible[:,np.newaxis,np.newaxis,:,np.newaxis]) ) unnorm_s1 = unnorm_s1 * considered_signals[np.newaxis,np.newaxis,:,:,np.newaxis] s1 = unnorm_s1 / unnorm_s1.sum(axis=-2, keepdims=True) s1 = T.switch(T.isnan(s1), 0., s1) l2 = s1 / s1.sum(axis=-1, keepdims=True) expected_dist_l2 = T.tensordot(l2, distances, axes=[[4],[0]]) unnorm_l2 = T.exp(choice_alphas[:,np.newaxis,np.newaxis,np.newaxis,np.newaxis]*-expected_dist_l2) shape = unnorm_l2.shape _, picsize_index, _, _, state_index = T.mgrid[ 0:shape[0], 0:shape[1], 0:shape[2], 0:shape[3], 0:shape[4] ] unnorm_l2 = T.switch(state_index > picsize_index, 0, unnorm_l2) l2 = unnorm_l2 / T.sum(unnorm_l2, axis=-1, keepdims=True) l2_language = l2[:,:,T.arange(real_signals_indices.shape[0]), real_signals_indices,:].squeeze() costs_real = T.outer(cost_factors, objective_costs_real) utility_l2 = T.log(l2_language) unnorm_s3 = T.exp( alphas[:,np.newaxis,np.newaxis, np.newaxis]* (utility_l2 - costs_real[:,np.newaxis,:,np.newaxis]) ) s3 = unnorm_s3 / unnorm_s3.sum(axis=-2, keepdims=True) s3 = T.switch(T.isnan(s3), 0, s3) return_value = T.jacobian(s3.flatten(), alphas) return tt.function([ possible_signals_array, real_signals_indices, alphas, choice_alphas, cost_factors, objective_costs_possible, at_most_as_costly, types, distances ], return_value, on_unused_input='warn' ) ``` So to calculate the Jacobian: ```python s3_function_grad = theano_RSA() ``` To evaluate the function at a point, create some toy data: ```python num_participants = 1 num_states = 2 possible_signals_array = np.array([ [[0, 0], [0, 0]], [[1, 0], [0, 1]] ]) real_signals_indices = np.array([0]) costs_possible = np.array([1, 1]) at_most_as_costly = np.array([ [1,1], [1,1] ]) types = np.array([1,1]) distances = np.array([ [0,1], [1,0] ]) picked_signals_indices = np.array([0]) picsizes_values = np.array([1]) participants_indices = np.array([0]) states_values = np.array([1]) alphas = np.array([1.]) choice_alphas = np.array([1.]) cost_factors = np.array([0.01]) ``` And finally, evaluate the gradient, which is all NaN: ```python s3_function_grad( possible_signals_array = possible_signals_array, real_signals_indices = real_signals_indices, alphas = alphas, choice_alphas = choice_alphas, cost_factors = cost_factors, objective_costs_possible = costs_possible, at_most_as_costly = at_most_as_costly, types = types, distances = distances ) ``` -- --- You received this message because you are subscribed to the Google Groups "theano-users" group. To unsubscribe from this group and stop receiving emails from it, send an email to theano-users+unsubscr...@googlegroups.com. To view this discussion on the web visit https://groups.google.com/d/msgid/theano-users/ae4e649b-0319-47d0-917e-7447c23ce541n%40googlegroups.com.