While-loop in JAX not differentiable? How to get diff-able while-loop?

后端未结

关注

 0  1235

In TensorFlow, the tf.while_loop is fully differentiable.

In JAX, according to the documentation of jax.lax.while_loop:

while_l


                      
              相关标签: