-
Going through some Jax source code, I see that a lot of functions use
|
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 1 reply
-
Other statements that can be used within What will not work within JIT is attempting to branch based on dynamic quantites, such as values within arrays. There's some discussion of these topics in How To Think In JAX. |
Beta Was this translation helpful? Give feedback.
jit
is compatible withif
statements whose truth or falsehood is decidable at trace time.x is None
fits this criterion, because it concerns the identity of the Python variable.Other statements that can be used within
if
statements are static attributes of the an array, such asshape
,size
, ordtype
.What will not work within JIT is attempting to branch based on dynamic quantites, such as values within arrays.
There's some discussion of these topics in How To Think In JAX.