Merge pull request #41762 from zhuzilin:jit-scope-doc

PiperOrigin-RevId: 326629726
Change-Id: Ica8b43b09b634e1e5c6509c834340deaeb07a02c

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