# Copyright 2017 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""State management for eager execution."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import collections

import contextlib

import copy

import random

import threading

import numpy as np

import six

from tensorflow.core.protobuf import config_pb2

from tensorflow.core.protobuf import rewriter_config_pb2

from tensorflow.python import pywrap_tensorflow

from tensorflow.python import tf2

from tensorflow.python.eager import executor

from tensorflow.python.eager import monitoring

from tensorflow.python.framework import c_api_util

from tensorflow.python.framework import device as pydev

from tensorflow.python.util import compat

from tensorflow.python.util import is_in_graph_mode

from tensorflow.python.util import tf_contextlib

from tensorflow.python.util.tf_export import tf_export

GRAPH_MODE = 0

EAGER_MODE = 1

default_execution_mode = EAGER_MODE if tf2.enabled() else GRAPH_MODE

# Cache from (old_device_name, partial_new_device_name) -> (new_device_name,

# new_device_spec).

# Note that we do not protect this with a lock and instead rely on python's GIL

# and the idempotent nature of writes to provide thread safety.

_device_parsing_cache = {}

_starting_device_spec = pydev.DeviceSpec.from_string("")

_MAXINT32 = 2**31 - 1

DEVICE_PLACEMENT_EXPLICIT = pywrap_tensorflow.TFE_DEVICE_PLACEMENT_EXPLICIT

DEVICE_PLACEMENT_WARN = pywrap_tensorflow.TFE_DEVICE_PLACEMENT_WARN

DEVICE_PLACEMENT_SILENT = pywrap_tensorflow.TFE_DEVICE_PLACEMENT_SILENT

DEVICE_PLACEMENT_SILENT_FOR_INT32 = (

pywrap_tensorflow.TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32)

SYNC = 0

ASYNC = 1

MIRRORING_NONE = pywrap_tensorflow.TFE_MIRRORING_NONE

MIRRORING_ALL = pywrap_tensorflow.TFE_MIRRORING_ALL

_tf2_gauge = monitoring.BoolGauge("/tensorflow/api/tf2_enable",

"Whether tf2.enable() is called.")

_python_eager_context_create_counter = monitoring.Counter(

"/tensorflow/api/python/eager_context_create_counter",

"Counter for number of eager contexts created in Python.")

_tf2_gauge.get_cell().set(tf2.enabled())

class _EagerTensorCache(object):

"""Simple cache which evicts items based on length in a FIFO manner."""

def __init__(self, max_items=256, max_tensor_size=10000):

self._data = collections.OrderedDict()

self._max_items = max_items

self._max_tensor_size = max_tensor_size

def put(self, key, value):

if value._num_elements() > self._max_tensor_size: # pylint: disable=protected-access

return

self._data[key] = value

if len(self._data) > self._max_items:

self._data.popitem(last=False)

def get(self, key):

return self._data.get(key, None)

def flush(self):

self._data = {}

class FunctionCallOptions(object):

"""Options applied at call sites of eager functions.

Eager functions are functions decorated with tf.contrib.eager.defun.

"""

def __init__(self, executor_type=None, config_proto=None):

"""Constructor.

Args:

executor_type: (optional) name of the executor to be used to execute the

eager function. If None or an empty string, the default Tensorflow

executor will be used.

config_proto: (optional) a `config_pb2.ConfigProto` proto or

a serialized string of that proto.

The config used by Grappler when optimizing the function graph.

Each concrete function is optimized the first time is called. Changing

config_proto after the first call has no effect.

If config_proto is None, an empty RewriterConfig will be used.

"""

self.config_proto_serialized = config_proto

self.executor_type = executor_type

@property

def executor_type(self):

return self._executor_type

@executor_type.setter

def executor_type(self, executor_type):

self._executor_type = executor_type

@property

def config_proto_serialized(self):

return self._config_proto_serialized

@config_proto_serialized.setter

def config_proto_serialized(self, config):

if isinstance(config, config_pb2.ConfigProto):

self._config_proto_serialized = config.SerializeToString()

elif isinstance(config, str):

self._config_proto_serialized = config

elif config is None:

self._config_proto_serialized = (

config_pb2.ConfigProto().SerializeToString())

else:

raise ValueError("the rewriter config must be either a "

"config_pb2.ConfigProto, or a serialized string of that "

"proto or None. got: {}".format(type(config)))

# Map from context_id (an int) to _TensorCaches.

# Dicts are thread safe in CPython.

# TODO(iga): Remove this once TensorCaches are moved to C++.

_tensor_caches_map = {}

class _TensorCaches(threading.local):

"""Thread local tensor caches."""

def __init__(self):

super(_TensorCaches, self).__init__()

self._ones_rank_cache = None

self._zeros_cache = None

@property

def ones_rank_cache(self):

if not self._ones_rank_cache:

self._ones_rank_cache = _EagerTensorCache()

return self._ones_rank_cache

@property

def zeros_cache(self):

if not self._zeros_cache:

self._zeros_cache = _EagerTensorCache()

return self._zeros_cache

class _ThreadLocalData(threading.local):

"""Thread local storage for the eager context."""

def __init__(self):

super(_ThreadLocalData, self).__init__()

self.device_spec = _starting_device_spec

self.device_name = ""

self.mode = default_execution_mode

self.is_eager = default_execution_mode == EAGER_MODE

self.scope_name = ""

self.summary_writer = None

self.summary_recording = None

self.summary_recording_distribution_strategy = True

self.summary_step = None

self.function_call_options = None

self.executor = None

ContextSwitch = collections.namedtuple(

"ContextSwitch", ["is_building_function", "enter_context_fn",

"device_stack"])

# `_ContextSwitchStack` is a `threading.local` to match the semantics of

# ``DefaultGraphStack`, which is also a `threading.local`.

class _ContextSwitchStack(threading.local):

"""A thread-local stack of context switches."""

def __init__(self, eager):

super(_ContextSwitchStack, self).__init__()

self.stack = []

if eager:

# Initialize the stack with a pointer to enter the eager context; this

# ensures that the fact that eager execution was enabled is propagated

# across threads, since (1) `enable_eager_execution` modifies a

# process-level flag (`default_execution_mode`) and (2) `__init__` is

# called each time a threading.local object is used in a separate thread.

self.push(is_building_function=False, enter_context_fn=eager_mode,

device_stack=None)

def push(self, is_building_function, enter_context_fn, device_stack):

"""Push metadata about a context switch onto the stack.

A context switch can take any one of the two forms: installing a graph as

the default graph, or entering the eager context. For each context switch,

we record whether or not the entered context is building a function.

Args:

is_building_function: (bool.) Whether the context is building a function.

enter_context_fn: (function.) A callable that executes the context switch.

For example, `graph.as_default` or `eager_mode`.

device_stack: If applicable, the device function stack for this

graph. When breaking out of graphs in init_scope, the innermost nonempty

device stack is used. Eager contexts put `None` here and the value is

never used.

"""

self.stack.append(

ContextSwitch(is_building_function, enter_context_fn, device_stack))

def pop(self):

"""Pop the stack."""

self.stack.pop()

class LogicalDevice(

collections.namedtuple("LogicalDevice", ["name", "device_type"])):

"""Abstraction for a device initialized by the runtime.

A LogicalDevice corresponds to a initialized instance on a PhysicalDevice or a

remote device available in the cluster. Tensors and operations can be placed

on a specific LogicalDevice by calling `tf.device()` with the `name` of the

LogicalDevice.

Fields:

name: The fully qualified name of the device. Can be used for Op or function

placement.

device_type: String declaring the type of device such as "CPU" or "GPU".

"""

pass

@tf_export("config.experimental.VirtualDeviceConfiguration")

class VirtualDeviceConfiguration(

collections.namedtuple("VirtualDeviceConfiguration", ["memory_limit"])):

"""Configuration class for virtual devices for a PhysicalDevice.

Fields:

memory_limit: (optional) Maximum memory (in MB) to allocate on the virtual

device. Currently only supported for GPUs.

"""

def __new__(cls, memory_limit=None):

return super(VirtualDeviceConfiguration, cls).__new__(cls, memory_limit)

class PhysicalDevice(

collections.namedtuple("PhysicalDevice", ["name", "device_type"])):

"""Abstraction for a locally visible physical device.

TensorFlow can utilize various devices such as the CPU or multiple GPUs

for computation. Before initializing a local device for use, the user can

customize certain properties of the device such as it's visibility or memory

configuration.

Once a PhysicalDevice is initialized one or many LogicalDevice objects are

created. Use tf.config.set_virtual_device_configuration() to create multiple

LogicalDevice objects for a PhysicalDevice. This is useful when separation

between models is needed.

Fields:

name: Unique identifier for device.

device_type: String declaring the type of device such as "CPU" or "GPU".

"""

pass

class _AtomicCounter(object):

"""A simple atomic counter."""

def __init__(self):

self._value = 0

self._lock = threading.Lock()

def increment_and_get(self):

with self._lock:

self._value += 1

return self._value

_context_id_counter = _AtomicCounter()

class _TensorCacheDeleter(object):

"""Deletes tensor caches for a given context."""

def __init__(self, context_id):

self._context_id = context_id

def __del__(self):

if _tensor_caches_map is None:

return

if self._context_id in _tensor_caches_map:

del _tensor_caches_map[self._context_id]

# TODO(agarwal): rename to EagerContext / EagerRuntime ?

# TODO(agarwal): consider keeping the corresponding Graph here.

class Context(object):

"""Environment in which eager operations execute."""

# TODO(agarwal): create and link in some documentation for `execution_mode`.

# pylint: disable=redefined-outer-name

def __init__(self,

config=None,

device_policy=None,

execution_mode=None,

server_def=None):

"""Creates a new Context.

Args:

config: (Optional.) A `ConfigProto` protocol buffer with configuration

options for the Context. Note that a lot of these options may be

currently unimplemented or irrelevant when eager execution is enabled.

device_policy: (Optional.) What policy to use when trying to run an

operation on a device with inputs which are not on that device.

When set to None, an appropriate value will be picked automatically.

The value picked may change between TensorFlow releases.

Defaults to DEVICE_PLACEMENT_SILENT.

Valid values:

- DEVICE_PLACEMENT_EXPLICIT: raises an error if the placement is

not correct.

- DEVICE_PLACEMENT_WARN: copies the tensors which are not on the

right device but raises a warning.

- DEVICE_PLACEMENT_SILENT: silently copies the tensors. This might

hide performance problems.

- DEVICE_PLACEMENT_SILENT_FOR_INT32: silently copies int32 tensors,

raising errors on the other ones.

execution_mode: (Optional.) Policy controlling how operations dispatched

are actually executed. When set to None, an appropriate value will be

picked automatically. The value picked may change between TensorFlow

releases.

Valid values:

- SYNC: executes each operation synchronously.

- ASYNC: executes each operation asynchronously. These

operations may return "non-ready" handles.

server_def: (Optional.) A tensorflow::ServerDef proto.

Enables execution on remote devices. GrpcServers need to be started by

creating an identical server_def to this, and setting the appropriate

task_indexes, so that the servers can communicate. It will then be

possible to execute operations on remote devices.

Raises:

ValueError: If execution_mode is not valid.

"""

# This _id is used only to index the tensor caches.

# TODO(iga): Remove this when tensor caches are moved to C++.

self._id = _context_id_counter.increment_and_get()

self._tensor_cache_deleter = _TensorCacheDeleter(self._id)

_tensor_caches_map[self._id] = _TensorCaches()

self._config = config

self._thread_local_data = _ThreadLocalData()

self._context_switches = _ContextSwitchStack(self.executing_eagerly())

self._context_handle = None

self._context_devices = None

self._post_execution_callbacks = []

self._seed = None

self._initialize_lock = threading.Lock()

if device_policy is None:

device_policy = DEVICE_PLACEMENT_SILENT

self._device_policy = device_policy

self._mirroring_policy = None

if execution_mode not in (None, SYNC, ASYNC):

raise ValueError(

"execution_mode should be None/SYNC/ASYNC. Got %s" % execution_mode)

if execution_mode is None:

execution_mode = SYNC

self._default_is_async = execution_mode == ASYNC

self._server_def = server_def

self._collective_ops_server_def = None

self._collective_leader = None

self._collective_scoped_allocator_enabled_ops = None

self._collective_use_nccl_communication = None

self._collective_device_filters = None

self._device_lock = threading.Lock()

self._physical_devices = None

self._visible_device_list = []

self._memory_growth_map = None

self._virtual_device_map = {}

# Values set after construction

self._optimizer_jit = None

self._intra_op_parallelism_threads = None

self._inter_op_parallelism_threads = None

self._soft_device_placement = None

self._log_device_placement = None

self._optimizer_experimental_options = {}

_python_eager_context_create_counter.get_cell().increase_by(1)

# pylint: enable=redefined-outer-name

def _set_global_seed(self, seed):

"""Set a global eager mode seed for random ops."""

self._seed = seed

# `random.Random(seed)` needs `seed` to be hashable, while values of type

# e.g. `np.int64` or `np.ndarray` are not. We use `int(...)` to convert them

# to int.

try:

hash(seed)

except TypeError:

seed = int(np.array(seed))

self._rng = random.Random(seed)

# Also clear the kernel cache, to reset any existing seeds

if self._context_handle is not None:

pywrap_tensorflow.TFE_ContextClearCaches(self._context_handle)

def _internal_operation_seed(self):

"""Returns a fake operation seed.

In eager mode, user shouldn't set or depend on operation seed.

Here, we generate a random seed based on global seed to make

operation's randomness different and depend on the global seed.

Returns:

A fake operation seed based on global seed.

"""

return self._rng.randint(0, _MAXINT32)

def _initialize_logical_devices(self):

"""Helper to initialize devices."""

# Store list of devices

self._logical_devices = []

self._context_devices = []

device_list = pywrap_tensorflow.TFE_ContextListDevices(

self._context_handle)

try:

self._num_gpus = 0

for i in range(pywrap_tensorflow.TF_DeviceListCount(device_list)):

dev_name = pywrap_tensorflow.TF_DeviceListName(device_list, i)

self._context_devices.append(pydev.canonical_name(dev_name))

spec = pydev.DeviceSpec.from_string(dev_name)

self._logical_devices.append(

LogicalDevice(name=dev_name, device_type=spec.device_type))

dev_type = pywrap_tensorflow.TF_DeviceListType(device_list, i)

if dev_type == "GPU":

self._num_gpus += 1

finally:

pywrap_tensorflow.TF_DeleteDeviceList(device_list)

def ensure_initialized(self):

"""Initialize handle and devices if not already done so."""

with self._initialize_lock:

if self._context_handle is not None:

return

assert self._context_devices is None

opts = pywrap_tensorflow.TFE_NewContextOptions()

try:

config_str = self.config.SerializeToString()

pywrap_tensorflow.TFE_ContextOptionsSetConfig(opts, config_str)

if self._device_policy is not None:

pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy(

opts, self._device_policy)

if self._mirroring_policy is not None:

pywrap_tensorflow.TFE_ContextOptionsSetMirroringPolicy(

opts, self._mirroring_policy)

if self._default_is_async == ASYNC:

pywrap_tensorflow.TFE_ContextOptionsSetAsync(opts, True)

self._context_handle = pywrap_tensorflow.TFE_NewContext(opts)

finally:

pywrap_tensorflow.TFE_DeleteContextOptions(opts)

assert not (self._server_def and self._collective_ops_server_def), (

"Cannot enable remote execution as well as collective ops at the "

"moment. If this is important to you, please file an issue.")

if self._server_def is not None:

server_def_str = self._server_def.SerializeToString()

pywrap_tensorflow.TFE_ContextSetServerDef(self._context_handle, 600,

server_def_str)

elif self._collective_ops_server_def is not None:

server_def_str = self._collective_ops_server_def.SerializeToString()

pywrap_tensorflow.TFE_EnableCollectiveOps(self._context_handle,

server_def_str)

self._initialize_logical_devices()

def _clear_caches(self):

self.ones_rank_cache().flush()

self.zeros_cache().flush()

pywrap_tensorflow.TFE_ClearScalarCache()

def set_server_def(self, server_def, keep_alive_secs=600):

"""Allow setting a server_def on the context.

When a server def is replaced, it effectively clears a bunch of caches

within the context. If you attempt to use a tensor object that was pointing

to a tensor on the remote device, it will raise an error.

Args:

server_def: A tensorflow::ServerDef proto.

Enables execution on remote devices.

keep_alive_secs: Num. seconds after which the remote end will hang up.

As long as the client is still alive, the server state for the context

will be kept alive. If the client is killed (or there is some failure),

the server will clean up its context keep_alive_secs after the final RPC

it receives.

Raises:

ValueError: if server_def is None.

"""

if not server_def:

raise ValueError("server_def is None.")

self._server_def = server_def

if self._context_handle:

server_def_str = server_def.SerializeToString()

pywrap_tensorflow.TFE_ContextSetServerDef(self._context_handle,

keep_alive_secs, server_def_str)

self._initialize_logical_devices()

# Clear all the caches in case there are remote tensors in them.

self._clear_caches()

def enable_collective_ops(self, server_def):

"""Enable distributed collective ops with an appropriate server_def.

Args:

server_def: A tensorflow::ServerDef proto. Enables execution on remote

devices.

Raises:

ValueError: if server_def is None.

RuntimeError: if this method is not called at program startup.

"""

if not server_def:

raise ValueError("server_def is None.")

if self._context_handle is not None:

raise RuntimeError("Collective ops must be enabled at program startup")

self._collective_ops_server_def = server_def

def configure_collective_ops(

self,

collective_leader="",

scoped_allocator_enabled_ops=("CollectiveReduce",),

use_nccl_communication=False,

device_filters=None):

"""Configure collective ops.

Collective group leader is necessary for collective ops to run, other

configurations are mainly for the purpose of performance.

Args:

collective_leader: a device string for collective leader, e.g.

"/job:worker/replica:0/task:"; empty string means local execution of

collective ops.

scoped_allocator_enabled_ops: a tuple or a list of op names for scoped

allocator to run with.

use_nccl_communication: whether to use nccl communication for collective

ops.

device_filters: a tuple or a list of device strings. If set, corresponding

task can only see the devices filtered by these device filters.

Raises:

RuntimeError: if this method is not called at program startup.

"""

if self._collective_leader is not None:

if (self._collective_leader != collective_leader or

self._collective_scoped_allocator_enabled_ops !=

scoped_allocator_enabled_ops or

self._collective_use_nccl_communication != use_nccl_communication or

self._collective_device_filters != device_filters):

raise ValueError("Collective ops are already configured.")

else:

return

if self._context_handle is not None:

raise RuntimeError("Collective ops must be configured at program startup")

self._collective_leader = collective_leader

self._collective_scoped_allocator_enabled_ops = scoped_allocator_enabled_ops

self._collective_use_nccl_communication = use_nccl_communication

self._collective_device_filters = device_filters

@property

def _handle(self):

if self._context_handle is None:

raise AssertionError("Context must be initialized first.")

return self._context_handle

@property

def _devices(self):

if self._context_devices is None:

raise AssertionError("Context must be initialized first.")

return self._context_devices

def __str__(self):

if self._context_handle is None:

return "Eager TensorFlow Context. Devices currently uninitialized."

else:

devices = self._devices

lines = ["Eager TensorFlow Context with %d devices" % (len(devices))]

for i, d in enumerate(devices):

lines.append(" Device %d: %s" % (i, d))

return "\n".join(lines)

@tf_contextlib.contextmanager

def _mode(self, mode):

"""A context manager to allow setting the mode to EAGER/GRAPH."""

ctx = self._thread_local_data

old_mode = ctx.mode

old_is_eager = ctx.is_eager

ctx.mode = mode

ctx.is_eager = mode == EAGER_MODE

if mode == EAGER_MODE:

# Entering graph mode does not provide us with sufficient information to

# record a context switch; graph-based context switches are only logged

# when a graph is registered as the default graph.

self.context_switches.push(False, eager_mode, None)

try:

yield

finally:

ctx.is_eager = old_is_eager

ctx.mode = old_mode

if mode == EAGER_MODE:

self.context_switches.pop()

def executing_eagerly(self):

"""Returns True if current thread has eager executing enabled."""

return self._thread_local_data.is_eager

def ones_rank_cache(self):

"""Per-device cache for scalars."""

return _tensor_caches_map[self._id].ones_rank_cache

def zeros_cache(self):

"""Per-device cache for scalars."""

return _tensor_caches_map[self._id].zeros_cache

@property

def scope_name(self):

"""Returns scope name for the current thread."""

return self._thread_local_data.scope_name

@scope_name.setter

def scope_name(self, s):

"""Sets scope name for the current thread."""

self._thread_local_data.scope_name = s

@property

def summary_writer(self):

"""Returns default summary writer for the current thread."""

return self._thread_local_data.summary_writer

@summary_writer.setter

def summary_writer(self, writer):

"""Sets default summary writer for the current thread."""

self._thread_local_data.summary_writer = writer

@property

def summary_recording(self):

"""Returns summary recording condition."""

return self._thread_local_data.summary_recording

@summary_recording.setter

def summary_recording(self, condition):

"""Sets summary recording condition."""

self._thread_local_data.summary_recording = condition

@property

def summary_recording_distribution_strategy(self):

"""Returns summary recording condition for distribution strategy."""

return self._thread_local_data.summary_recording_distribution_strategy

@summary_recording_distribution_strategy.setter

def summary_recording_distribution_strategy(self, condition):

"""Sets summary recording condition for distribution strategy."""

self._thread_local_data.summary_recording_distribution_strategy = condition

@property

def summary_step(self):

"""Returns summary step variable."""

return self._thread_local_data.summary_step

@summary_step.setter

def summary_step(self, step):

"""Sets summary step variable."""

self._thread_local_data.summary_step = step

@property

def device_name(self):

"""Returns the device name for the current thread."""

return self._thread_local_data.device_name

@property

def device_spec(self):

"""Returns the device spec for the current thread."""

return self._thread_local_data.device_spec

def _set_device(self, device_name, device_spec):

self._thread_local_data.device_name = device_name

self._thread_local_data.device_spec = device_spec

def device(self, name):

"""Context-manager to force placement of operations and Tensors on a device.

Args:

name: Name of the device or None to get default placement.

Returns:

Context manager that forces device placement.

Raises:

ValueError: If name is not a string or is an invalid device name.

RuntimeError: If device scopes are not properly nested.

"""

return _EagerDeviceContext(self, name)

def devices(self):

"""List of the names of devices available to execute operations."""

return self._devices

# TODO(fishx): remove this property.

@property

def execution_mode(self):

"""Gets execution mode for current thread."""

return ASYNC if self.is_async() else SYNC

@execution_mode.setter

def execution_mode(self, mode):

"""Sets execution mode for current thread."""

if mode not in (None, SYNC, ASYNC):

raise ValueError(

"Execution mode should be None/SYNC/ASYNC. Got %s" % mode)

if mode is None:

mode = SYNC

enable_async = (mode == ASYNC)

if self.is_async() != enable_async:

# Only set the execution mode if the context has already been initialized

if self._context_handle is not None:

self.async_wait()

executor_new = executor.Executor(enable_async)

self._thread_local_data.executor = executor_new

pywrap_tensorflow.TFE_ContextSetExecutorForThread(

self._context_handle, executor_new.handle())

else:

self._default_is_async = enable_async

def is_async(self):

if self._thread_local_data.executor is None:

return self._default_is_async

else:

return self._thread_local_data.executor.is_async()

@property

def executor(self):

return self._thread_local_data.executor

@executor.setter

def executor(self, e):

ensure_initialized()

if self._thread_local_data.executor != e:

self._thread_local_data.executor = e

if e is None:

pywrap_tensorflow.TFE_ContextClearExecutorForThread(

self._context_handle)

else:

pywrap_tensorflow.TFE_ContextSetExecutorForThread(

self._context_handle, e.handle())

@property

def config(self):

"""Return the ConfigProto with all runtime deltas applied."""

# Ensure physical devices have been discovered and config has been imported

self._initialize_physical_devices()

config = config_pb2.ConfigProto()

if self._config is not None:

config.CopyFrom(self._config)

if self._optimizer_jit is not None:

config.graph_options.optimizer_options.global_jit_level = (

config_pb2.OptimizerOptions.ON_1

if self._optimizer_jit else config_pb2.OptimizerOptions.OFF)

if self._intra_op_parallelism_threads is not None:

config.intra_op_parallelism_threads = self._intra_op_parallelism_threads

if self._inter_op_parallelism_threads is not None:

config.inter_op_parallelism_threads = self._inter_op_parallelism_threads

if self._soft_device_placement is not None:

config.allow_soft_placement = self._soft_device_placement

else:

config.allow_soft_placement = self.executing_eagerly()

if self._log_device_placement is not None:

config.log_device_placement = self._log_device_placement

def rewriter_toggle(option):

toggle = self._optimizer_experimental_options.get(option, None)

if toggle is None:

return

setattr(config.graph_options.rewrite_options,

option,

(rewriter_config_pb2.RewriterConfig.ON

if toggle else rewriter_config_pb2.RewriterConfig.OFF))

def rewriter_bool(option):

toggle = self._optimizer_experimental_options.get(option, None)

if toggle is None:

return

setattr(config.graph_options.rewrite_options,

option,

toggle)

rewriter_toggle("layout_optimizer")

rewriter_toggle("constant_folding")

rewriter_toggle("shape_optimization")

rewriter_toggle("remapping")

rewriter_toggle("arithmetic_optimization")

rewriter_toggle("dependency_optimization")

rewriter_toggle("loop_optimization")

rewriter_toggle("function_optimization")

rewriter_toggle("debug_stripper")

rewriter_bool("disable_model_pruning")

rewriter_toggle("scoped_allocator_optimization")

rewriter_toggle("pin_to_host_optimization")

rewriter_toggle("implementation_selector")

rewriter_toggle("auto_mixed_precision")

rewriter_bool("disable_meta_optimizer")

nodes = self._optimizer_experimental_options.get("min_graph_nodes", None)

if nodes is not None:

config.graph_options.rewrite_options.min_graph_nodes = nodes

# Compute device counts

config.device_count["CPU"] = 0

config.device_count["GPU"] = 0

for dev in self._physical_devices:

if dev not in self._visible_device_list:

continue

virtual_devices = self._virtual_device_map.get(dev)

if virtual_devices is None:

config.device_count[dev.device_type] += 1

else:

config.device_count[dev.device_type] += len(virtual_devices)

# Configure gpu_options

gpu_options = self._compute_gpu_options()

config.gpu_options.MergeFrom(gpu_options)

# Configure collective ops

if self._collective_leader:

config.experimental.collective_group_leader = self._collective_leader

if self._collective_scoped_allocator_enabled_ops:

rewrite_options = config.graph_options.rewrite_options

rewrite_options.scoped_allocator_optimization = (

rewriter_config_pb2.RewriterConfig.ON)

del rewrite_options.scoped_allocator_opts.enable_op[:]

for op in self._collective_scoped_allocator_enabled_ops:

rewrite_options.scoped_allocator_opts.enable_op.append(op)

if self._collective_use_nccl_communication:

config.experimental.collective_nccl = True

if self._collective_device_filters:

del config.device_filters[:]

for f in self._collective_device_filters:

config.device_filters.append(f)

return config

def _compute_gpu_options(self):

"""Build the GPUOptions proto."""

visible_device_list = []

virtual_devices = []

gpu_index = -1

memory_growths = set()

for dev in self.list_physical_devices("GPU"):

gpu_index += 1

if dev not in self._visible_device_list:

continue

growth = self._memory_growth_map[dev]

memory_growths.add(growth)

visible_device_list.append(str(gpu_index))

if self._virtual_device_map:

vdevs = self._virtual_device_map.get(dev, [])

device_limits = []

for virt_dev in vdevs:

device_limits.append(virt_dev.memory_limit)

virtual_devices.append(

config_pb2.GPUOptions.Experimental.VirtualDevices(

memory_limit_mb=device_limits))

# Only compute growth if virtual devices have not been configured and we

# have GPUs

if not virtual_devices and memory_growths:

if len(memory_growths) > 1:

raise ValueError("Memory growth cannot differ between GPU devices")

allow_growth = memory_growths.pop()

else:

allow_growth = None

return config_pb2.GPUOptions(

allow_growth=allow_growth,

visible_device_list=",".join(visible_device_list),

experimental=config_pb2.GPUOptions.Experimental(

virtual_devices=virtual_devices))

@property

def function_call_options(self):

"""Returns function call options for current thread.

Note that the returned object is still referenced by the eager context.

Returns: the FunctionCallOptions for current thread.

"""

if self._thread_local_data.function_call_options is None:

config = self.config

# Default to soft placement for functions unless specified

if self._soft_device_placement is None:

config.allow_soft_placement = True

self._thread_local_data.function_call_options = FunctionCallOptions(

config_proto=config)

return self._thread_local_data.function_call_options

@function_call_options.setter

def function_call_options(self, options):

"""Returns function call options for current thread."""

self._thread_local_data.function_call_options = options

def async_wait(self):

"""Waits for ops dispatched in ASYNC mode to finish."""

pywrap_tensorflow.TFE_ContextAsyncWait(self._handle)

def async_clear_error(self):

"""Clears errors raised during ASYNC execution."""

pywrap_tensorflow.TFE_ContextAsyncClearError(self._handle)

def num_gpus(self):

"""The number of GPUs available to execute operations."""

self.ensure_initialized()

return self._num_gpus

def add_function(self, fn):

"""Add a function definition to the context.

Once added, the function (identified by its name) can be executed like any

other operation.

Args:

fn: A wrapped TF_Function (returned from TF_GraphToFunction_wrapper).

"""

self.ensure_initialized()

pywrap_tensorflow.TFE_ContextAddFunction(self._handle, fn)

def add_function_def(self, fdef):

"""Add a function definition to the context.

Once added, the function (identified by its name) can be executed like any

other operation.

Args:

fdef: A FunctionDef protocol buffer message.

"""

self.ensure_initialized()