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