Use of jax.lax.stop_gradient in nn.softmax

[amishra791]

Hi, I was taking a look at the softmax function and it's deprecated version. I noticed that in the deprecated version has jax.lax.stop_gradient. What is the use of this function here, and why is it not used in the most recent implementation of softmax?

[jakevdp]

The reason for this is that the non-deprecated softmax has a custom JVP, and so we don't differentiate through the implementation, which makes the stop_gradient call unnecessary.

[amishra791]

Thanks! Could you explain more as to why we needed the stop_gradient in the deprecated version?

[jakevdp]

See https://stackoverflow.com/questions/74822097/purpose-of-stop-gradient-in-jax-nn-softmax