tf.xla.experimental.jit_scope

Включает или отключает JIT-компиляцию операторов в области действия.

@contextlib.contextmanager
tf.xla.experimental.jit_scope(
    compile_ops=True, separate_compiled_gradients=False
)

Примечание: Это экспериментальная функция.

Компиляция служит подсказкой и поддерживается только по возможности.

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.

Пример 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')

Операции, не входящие в область действия, могут быть сгруппированы и скомпилированы вместе с операциями в области действия с помощью compile_ops=True, тогда как операции в области действия с compile_ops=False никогда не будут скомпилированы.

# 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

Если требуется скомпилировать только операции в области действия с compile_ops=True, рассмотрите возможность добавления внешнего 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