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')

Возвращаемое значение:

Текущая область, включающая или отключающая компиляцию.