✳flâneur — a map of the web's best reading

JAX Frequently Asked Questions (FAQ) — JAX documentation

jax.readthedocs.io · 2,250 words · saved by 1 readers

If you have a Python function that changes behavior after using jax.jit(), perhaps your function uses global state, or has side-effects. In the following code, the impure_func uses the global y and has a side-effect due to print: Without jit the output is: and with jit it is: For jax.jit(), the function is executed once using the Python interpreter, at which time the Inside printing happens, and the first value of y is observed. Then, the function is compiled and cached, and executed multiple times with different values of x, but with the same first value of y. Additional reading: JAX - The Sharp Bits Sometimes users are surprised by the fact that wrapping a function with jit() can change the function’s outputs. For example: This slight difference in output comes from optimizations within the XLA compiler: during compilation, XLA will sometimes rearrange or elide certain operations to make the overall computation more efficient. In this case, XLA utilizes the properties of the logarith

Frequently asked questions (FAQ) # We are collecting answers to frequently asked questions here. Contributions welcome! jit changes the behavior of my function # If you have a Python function that changes behavior after using jax.jit() , perhaps your function uses global state, or has side-effects. In the following code, the impure_func uses the global y and has a side-effect due to print : y = 0 # @jit # Different behavior with jit def impure_func ( x ): print ( "Inside:" , y ) return x + y for y in range ( 3 ): print ( "Result:" , impure_func ( y )) Without jit the output is: Inside : 0 Resu

Explore this link on the map →

related reading