View source on GitHub |
Enable or disable JIT compilation of operators within the scope.
@contextlib.contextmanager
tf.xla.experimental.jit_scope( compile_ops=True, separate_compiled_gradients=False )
The compilation is a hint and only supported on a best-effort basis.
Example usage | |
---|---|
|
Example of separate_compiled_gradients
:
# In the example below, the computations for f, g and h will all be compiled
# in separate scopes.
with tf.xla.experimental.jit_scope(
separate_compiled_gradients=True):
f = tf.matmul(a, b)
g = tf.gradients([f], [a, b], name='mygrads1')
h = tf.gradients([f], [a, b], name='mygrads2')
Ops that are not in the scope may be clustered and compiled with ops in
the scope with compile_ops=True
, while the ops in the scope with
compile_ops=False
will never be compiled.
For example | |
---|---|
|
If you want to only compile the ops in the scope with compile_ops=True
,
consider adding an outer jit_scope(compile_ops=False)
:
# In the example below, only x will be compiled.
with tf.xla.experimental.jit_scope(compile_ops=False):
with tf.xla.experimental.jit_scope():
x = tf.matmul(a, b)
y = tf.matmul(c, d)
loss = x + y
Raises | |
---|---|
RuntimeError
|
if called when eager execution is enabled. |
Yields | |
---|---|
The current scope, enabling or disabling compilation. |