Differenting jax.lax.fori_loop

[lyg1597]

Is it possible to compute the gradient of jax.lax.fori_loop with respect to upper? More specifically, consider function

def func(ub):
   def body_fun(i, val):
        return val+1
    res = jax.lax.fori_loop(0, ub, body_fun, 0)
    return res

where the result res will always be equal to ub. Is it possible in jax to differentiate this function with respect to it's input? Or maybe there exist some function other than fori_loop that I can potentially use to achieve this? Thank you very much in advance.

[jakevdp]

Differentiation is generally only well-defined with respect to floating point values, and since the upper bound of a for loop is integer typed, there is no well-defined gradient.

[lyg1597]

This makes sense. Thank you very much for the explaination.