I am confused with what types of operations are supported for automatic differentiation in tf. Concretely, is tensor indexing operation as follows supported?
...
# feat is output from some conv layer and the shape is B*H*W*C
# case one
loss = feat[:,1:,1:,:] - feat[:,:-1,:-1,:]
# case two
feat[:,1:,1:,:] = feat[:,1:,1:,:]/2. # assign and replace part original value
loss = tf.reduce_sum(feat)
This isn'ta direct answer, but as a clue, this automatic differentiation library autograd lists operations that are not supported, see Non-differentiable functions, for example floor(), round() are not auto differentiable.
One can also define their own operations, provided if you can code the gradients yourself, see extend-autograd-by-defining-your-own
I would guess tf is very similar to this.