Gradient is NaN
fausto....@gmail.com
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
)
```