Jax: detach() or stop_gradient() equivalent?

Created on 18 Jan 2020  路  2Comments  路  Source: google/jax

If we want certain jax array to be invariant to gradient operator, what is the jax equivalent that we can use (in PyTorch detach() function or in tensorflow stop_gradient() operator works)?

question
Source
picture anirbanl

Most helpful comment

jax.lax.stop_gradient does this: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.stop_gradient.html#jax.lax.stop_gradient

picture shoyer on 18 Jan 2020
🎉1 👍1

All 2 comments

jax.lax.stop_gradient does this: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.stop_gradient.html#jax.lax.stop_gradient

shoyer picture shoyer on 18 Jan 2020
🎉1 👍1

Thanks!

anirbanl picture anirbanl on 18 Jan 2020
Was this page helpful?
0 / 5 - 0 ratings

Related issues

alexbw picture alexbw  路  3Comments
rdaems picture rdaems  路  3Comments
kunc picture kunc  路  3Comments
lhk picture lhk  路  3Comments
harshit-2115 picture harshit-2115  路  3Comments