How is it possible that map_coordinates function is differentiable

[jakubMitura14]

I had used the map_coordinates function in the definition of loss function of the neural network and it worked - hence I assume that it has to be differentiable (I used Adam optimizer and loss value was steadily dropping).
However now I am analyzing the code https://jax.readthedocs.io/en/latest/_modules/jax/_src/scipy/ndimage.html#map_coordinates
, and I do not see how it can be differentiable - for example in function

_linear_indices_and_weights

the very first line is

lower = jnp.floor(coordinate)

is not differentiable and even if it would be we are acting here on coordinated to the array (looking for coordinates close by) and this is also not differentiable action, so How is it possible that it works?

Thanks for help !

[jakevdp]

map_coordinates is differentiable because it is built from differentiable components: for example, you mention jnp.floor, which itself is differentiable:

print(jax.grad(jnp.floor)(0.5))
# 0.0

The sense in which this is differentiable is that if you vary the input infinitesimally, the output does not change, so by the definition of the derivative, the gradient is zero. Things get a bit more complicated when jnp.floor is differentiated for integer inputs, but it still works if you take the derivative in the positive direction (which the floor gradient rule implicitly does implicitly).

Does that answer your question?

[jakubMitura14]

Thanks for response! ok about floor I assumed wrongly that lack of gradient variability of a function willl make optimazation impossible. Still what about indexing - this floor operation is calculating integers that act as indicies to the array - so we have discrete input and discrete set of possible output values (all in array). Hence gradient in my understanding should be zero or infinite (I may be wrong) and stability of differentiation through this function in theory should be extremely low.
As I stated at the very beginning it is working so sth is wrong with my reasoning

Thanks for help!

[jakevdp]

The gradient with respect to an indexing operation where the index is a float converted to an int is similar to the gradient with respect to floor: it's always zero (because infinitesimally changing the index cannot change the indexed value):

import jax
import jax.numpy as jnp

def f(x, i):
  return x[i.astype(int)]

x = jnp.arange(4.0)
i = 1.0
print(jax.grad(f, argnums=1)(x, i))
# 0.0