tf.xla.experimental.jit_scope

TensorFlow 1 version

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:

with tf.xla.experimental.jit_scope():
  c = tf.matmul(a, b)  # compiled
with tf.xla.experimental.jit_scope(compile_ops=False):
  d = tf.matmul(a, c)  # not compiled
with tf.xla.experimental.jit_scope(
    compile_ops=lambda node_def: 'matmul' in node_def.op.lower()):
  e = tf.matmul(a, b) + d  # matmul is compiled, the addition is not.

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:

# 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):

  # 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