Skip to content

Commit

Permalink
DOC: more info on disable_jit
Browse files Browse the repository at this point in the history
  • Loading branch information
jakevdp committed Mar 17, 2022
1 parent 1d5833d commit 625a69d
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion jax/_src/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -544,7 +544,12 @@ def disable_jit():
"""Context manager that disables :py:func:`jit` behavior under its dynamic context.
For debugging it is useful to have a mechanism that disables :py:func:`jit`
everywhere in a dynamic context.
everywhere in a dynamic context. Note that this not only disables explicit uses
of `jit` by the user, but will also remove any implicit JIT compilation used by the
JAX library: this includes implicit JIT computation of `body` and `cond`
functions passed to higher-level primitives like :func:`scan` and :func:`while_loop`,
JIT used in implementations of :mod:`jax.numpy` functions, and any other case where
`jit` is used within an API's implementation.
Values that have a data dependence on the arguments to a jitted function are
traced and abstracted. For example, an abstract value may be a
Expand Down

0 comments on commit 625a69d

Please sign in to comment.