Loss function seems to be miscalculated

Thank you for your implementation of IWAE. Although I have one remark and one question about it:
### The remark:
Your training loss function in ```iwae.py``` lines 71-73 seems miscalculated to me:
```python
loss = - \
tf.reduce_mean(tf.reduce_sum(
normalized_weights * log_weights, 0))
```
There seems to be a confusion between the loss function and the gradient estimate given in the paper. You seem to be trying to apply the trick in the paper that uses the normalization weights to keep gradient calculations parallel w.r.t the k samples of z despite the log. However, if so, you should stop the gradient at ```normalized_weights``` because they are only meant to weight the gradients, and not to be differentiated. Here is the gradient estimate I am referring to :

### The question :
I can see that, while calculating the weights you substract the maximum weight w.r.t the k samples of z as follows:
```python
weights = tf.exp(log_weights - tf.reduce_max(log_weights, 0))
```
I can also see that you somehow "put them back" after applying the average log when testing. I have seen a similar trick in the original implementation. I tried googling for the reason behind this trick using various keywords but I can't seem to find it. Could you please give me a quick explanation for this or direct me towards a resource that explains it.

Thanks again for your hard work on this implementation, as well as for your answers to my questions.