From 524200641643bc76e53939f25082726c7a98b5a7 Mon Sep 17 00:00:00 2001 From: zilinzhu Date: Mon, 27 Jul 2020 13:59:12 +0800 Subject: [PATCH] enrich docstring for jit_scope --- tensorflow/python/compiler/xla/jit.py | 28 +++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) 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