diff --git a/tensorflow/python/compiler/xla/jit.py b/tensorflow/python/compiler/xla/jit.py index 3ccf2959b76..fb911ec9637 100644 --- a/tensorflow/python/compiler/xla/jit.py +++ b/tensorflow/python/compiler/xla/jit.py @@ -70,6 +70,34 @@ def experimental_jit_scope(compile_ops=True, separate_compiled_gradients=False): 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: + + ```python + # In the example below, x and loss may be clustered and compiled together, + # while y will not be compiled. + with tf.xla.experimental.jit_scope(): + x = tf.matmul(a, b) + with tf.xla.experimental.jit_scope(compile_ops=False): + y = tf.matmul(c, d) + loss = x + y + ``` + + If you want to only compile the ops in the scope with `compile_ops=True`, + consider adding an outer `jit_scope(compile_ops=False)`: + + ```python + # 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 + ``` + Args: compile_ops: Whether to enable or disable compilation in the scope. Either a Python bool, or a callable that accepts the parameter