News for hackers. Stuff that matters.

A comment on JAX – NumPy on the CPU, GPU, and TPU, with great automatic differentiation. Parent comment

Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by cl3misch on Friday September 29, 2023 @09:29AM (#37703937)
This is still the case afaik.

For vanilla "if", the condition must be known at compile time. For runtime, you have to use "cond", "where", or "select" (which may be analogous).

[ Parent | on HN ]
Re:JAX – NumPy on the CPU, GPU, and TPU, with great automati... by mattjjatgoogle on Monday October 02, 2023 @01:28PM (#37741506)
Actually, that's never been a constraint for JAX autodiff. JAX grew out of the original Autograd (https://github.com/hips/autograd), so differentiating through Python control flow always worked. It's jax.jit and jax.vmap which place constraints on control flow, requiring structured control flow combinators like those.
[ Parent | on HN ]