def test_spmd_with_summary(self):
original_device_placement = config.get_soft_device_placement()
config.set_soft_device_placement(True)
strategy, _ = get_tpu_strategy(enable_spmd=True)
summary_dir = self.get_temp_dir()
writer = summary_ops.create_file_writer_v2(summary_dir)
with strategy.scope():
step = variables.Variable(0, dtype=dtypes.int64)
@def_function.function
def run():
with writer.as_default():
summary_ops.scalar("result", step * 2, step=step)
step.assign_add(1)
for _ in range(10):
strategy.run(run, args=())
for val in step.values:
for var in val.variables:
self.assertAllEqual(10, var)
config.set_soft_device_placement(original_device_placement)
def test_distribution(self, distribution):
input_array_1 = np.array([['a', 'b'], ['c', 'd']])
input_array_2 = np.array([['e', 'f'], ['g', 'h']])
inp_dataset = dataset_ops.DatasetV2.from_tensor_slices({
'input_1':
input_array_1,
'input_2':
input_array_2
})
inp_dataset = batch_wrapper(inp_dataset, 2, distribution)
# pyformat: disable
expected_output = [[b'a_X_e', b'a_X_f', b'b_X_e', b'b_X_f'],
[b'c_X_g', b'c_X_h', b'd_X_g', b'd_X_h']]
config.set_soft_device_placement(True)
with distribution.scope():
input_data_1 = keras.Input(shape=(2, ),
dtype=dtypes.string,
name='input_1')
input_data_2 = keras.Input(shape=(2, ),
dtype=dtypes.string,
name='input_2')
input_data = [input_data_1, input_data_2]
layer = category_crossing.CategoryCrossing()
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(inp_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_strategy_with_file(self, strategy):
# TODO(b/180614455): remove this check when MLIR bridge is always enabled.
if backend.is_tpu_strategy(strategy):
self.skipTest("This test needs MLIR bridge on TPU.")
vocab_data = ["earth", "wind", "and", "fire"]
vocab_file = self._write_to_temp_file("temp", vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = index_lookup.IndexLookup(
max_tokens=None,
num_oov_indices=1,
mask_token="",
oov_token="[OOV]",
dtype=dtypes.string,
vocabulary=vocab_file)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
model.compile(loss="mse")
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def testSoftPlacement(self):
if context.executing_eagerly():
self.assertTrue(config.get_soft_device_placement())
else:
self.assertFalse(config.get_soft_device_placement())
@def_function.function
def mod():
with ops.device('/device:GPU:0'):
a = constant_op.constant(1.0)
b = constant_op.constant(1.0)
return math_ops.mod(a, b)
config.set_soft_device_placement(True)
self.assertEqual(config.get_soft_device_placement(), True)
self.assertEqual(
config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is enabled, the mod operation should work with CPU
mod()
config.set_soft_device_placement(False)
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(
config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is disabled, the mod operation should fail on GPU
with self.assertRaises(errors.InvalidArgumentError):
mod()
def test_tpu_with_multiple_oov(self, strategy):
# TODO(b/180614455): remove this check when MLIR bridge is always enabled.
if "TPU" in type(strategy).__name__:
self.skipTest("This test needs MLIR bridge on TPU.")
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(
input_array).batch(2, drop_remainder=True)
expected_output = [[3, 4, 5, 6], [6, 5, 3, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = index_lookup.IndexLookup(max_tokens=None,
num_oov_indices=2,
mask_token="",
oov_token="[OOV]",
dtype=dtypes.string)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
model.compile(loss="mse")
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_distribution_strategy_output_with_adapt(self, strategy):
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(
input_array).batch(2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = text_vectorization.TextVectorization(
max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def setUp(self):
super(HardDevicePlacementTest, self).setUp()
context._reset_context()
config.set_soft_device_placement(enabled=False)
context.context().log_device_placement = True
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(context.context().soft_device_placement, False)
def setUp(self): super(ClusterPlacementTest, self).setUp() context._reset_context() config.set_soft_device_placement(enabled=True) context.context().log_device_placement = True workers, _ = test_util.create_local_cluster(2, 0) remote.connect_to_remote_host([workers[0].target, workers[1].target])
def test_distribution_strategy_output_with_adapt(self, strategy):
# TODO(b/180614455): remove this check when MLIR bridge is always enabled.
if backend.is_tpu_strategy(strategy):
self.skipTest("This test needs MLIR bridge on TPU.")
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(
input_array).batch(2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = text_vectorization.TextVectorization(
max_tokens=None,
standardize=None,
split=None,
output_mode=text_vectorization.INT)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_strategy(self, strategy):
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(
input_array).batch(2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = index_lookup.IndexLookup(max_tokens=None,
num_oov_indices=1,
mask_token="",
oov_token="[OOV]",
dtype=dtypes.string)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
model.compile(loss="mse")
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def DISABLED_test_tpu_distribution_with_file(self, distribution):
vocab_data = ["earth", "wind", "and", "fire"]
vocab_file = self._write_to_temp_file("temp", vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with distribution.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = index_lookup.IndexLookup(
max_tokens=None,
num_oov_indices=1,
mask_token="",
oov_token="[OOV]",
dtype=dtypes.string,
vocabulary=vocab_file)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
model.compile(loss="mse")
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_recover_from_compilation_failures(self, enable_packed_var):
# TODO(b/148150981): Stop skipping this test once recovery works
# for non-local TPU.
if FLAGS.tpu:
self.skipTest("Recovery fails for non-local TPU, see b/148150981")
# Disable automatic outside compilation.
config.set_soft_device_placement(False)
strategy = get_tpu_strategy(enable_packed_var)
@def_function.function
def compilation_failure_run():
def computation():
return random_ops.random_gamma([10], [0.5, 1.5])
return strategy.run(computation)
with self.assertRaises(errors.OpError):
compilation_failure_run()
@def_function.function
def good_run():
def computation():
return random_ops.random_normal([10])
return strategy.run(computation)
good_run()
def testSoftPlacement(self):
if context.executing_eagerly():
self.assertTrue(config.get_soft_device_placement())
else:
self.assertFalse(config.get_soft_device_placement())
def test_attr():
with ops.device('/device:GPU:0'):
return test_ops.test_attr(T=dtypes.float32, name='test_attr')
config.set_soft_device_placement(True)
self.assertEqual(config.get_soft_device_placement(), True)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is enabled, the test_attr operation should fallback
# to CPU with pure eager execution as well as functions
test_attr()
def_function.function(test_attr)()
config.set_soft_device_placement(False)
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is disabled, the test_attr operation should fail on
# GPU with pure eager execution as well as functions
with self.assertRaises(errors.InvalidArgumentError):
test_attr()
with self.assertRaises(errors.InvalidArgumentError):
def_function.function(test_attr)()
def testSoftPlacement(self):
if context.executing_eagerly():
self.assertTrue(config.get_soft_device_placement())
else:
self.assertFalse(config.get_soft_device_placement())
def mod():
with ops.device('/device:GPU:0'):
a = constant_op.constant(1.0)
b = constant_op.constant(1.0)
return math_ops.mod(a, b)
config.set_soft_device_placement(True)
self.assertEqual(config.get_soft_device_placement(), True)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is enabled, the mod operation should fallback to CPU
# with pure eager execution as well as functions
mod()
def_function.function(mod)()
config.set_soft_device_placement(False)
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is disabled, the mod operation should fail on GPU
# with pure eager execution as well as functions
with self.assertRaises(errors.InvalidArgumentError):
mod()
with self.assertRaises(errors.InvalidArgumentError):
def_function.function(mod)()
def testSoftPlacement(self):
self.assertEqual(config.get_soft_device_placement(), True)
@def_function.function
def mod():
with ops.device('/device:GPU:0'):
a = constant_op.constant(1.0)
b = constant_op.constant(1.0)
return math_ops.mod(a, b)
# Since soft placement is enabled, the mod operation should work with CPU
mod()
config.set_soft_device_placement(False)
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is disabled, the mod operation should fail on GPU
with self.assertRaises(errors.InvalidArgumentError):
mod()
config.set_soft_device_placement(True)
self.assertEqual(config.get_soft_device_placement(), True)
self.assertEqual(config.get_soft_device_placement(),
context.context().soft_device_placement)
# Since soft placement is re-enabled, the mod operation should work with CPU
mod()
def testGpuNone(self):
config.set_soft_device_placement(False)
gpus = config.list_physical_devices('GPU')
self.assertGreater(len(gpus), 0)
cpus = config.list_physical_devices('CPU')
self.assertEqual(len(cpus), 1)
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
self.assertGreater(len(config.get_visible_devices('GPU')), 0)
self.assertEqual(len(config.get_visible_devices('XLA_GPU')), 0)
config.set_visible_devices(cpus[0])
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
self.assertEqual(len(config.get_visible_devices('GPU')), 0)
self.assertEqual(len(config.list_logical_devices('XLA_GPU')), 0)
with self.assertRaisesRegex(errors.InvalidArgumentError,
'Could not satisfy'):
with ops.device('/device:GPU:0'):
a = array_ops.identity(1.0)
self.evaluate(a)
# Modifying the visible devices is not supported
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
config.set_visible_devices(gpus)
# Setting the same visible devices is fine
config.set_visible_devices(cpus[0])
def setUp(self): super(HardDevicePlacementTest, self).setUp() context._context = None ops.enable_eager_execution_internal() config.set_soft_device_placement(enabled=False) context.context().log_device_placement = True self.assertEqual(config.get_soft_device_placement(), False) self.assertEqual(context.context().soft_device_placement, False)
def setUp(self):
super(HardDevicePlacementTest, self).setUp()
context._reset_context()
config.set_soft_device_placement(enabled=False)
context.context().log_device_placement = True
cpus = context.context().list_physical_devices('CPU')
# Set 2 virtual CPUs
context.context().set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(context.context().soft_device_placement, False)
def testSoftPlacement(self):
# Temporarily replace the context
# pylint: disable=protected-access
old_context = context.context()
context._set_context(context.Context())
try:
config.set_device_policy('silent')
config.set_soft_device_placement(True)
cpu_tensor = constant_op.constant(1.0)
result = cpu_tensor + cpu_tensor
self.assertEqual(result.device,
'/job:localhost/replica:0/task:0/device:GPU:0')
finally:
context._set_context(old_context)
def test_distribution(self, distribution):
input_data = np.asarray([["omar"], ["stringer"], ["marlo"], ["wire"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_data).batch(
2, drop_remainder=True)
expected_output = [[0], [0], [1], [0]]
config.set_soft_device_placement(True)
with distribution.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = hashing.Hashing(num_bins=2)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def testSoftPlacement(self):
# Temporarily replace the context
# pylint: disable=protected-access
del context._context
context._context = context.Context()
try:
config.set_device_policy('silent')
config.set_soft_device_placement(True)
cpu_tensor = constant_op.constant(1.0)
result = cpu_tensor + cpu_tensor
self.assertEqual(result.device,
'/job:localhost/replica:0/task:0/device:GPU:0')
finally:
del context._context
context._context = context.Context()
def testSoftPlacement(self):
# Temporarily replace the context
# pylint: disable=protected-access
old_context = context.context()
context._set_context(context.Context())
try:
config.set_device_policy('silent')
config.set_soft_device_placement(True)
# Avoid the TensorHandle cache hit.
# TODO(b/169790439): include Context to the TensorHandle cache.
cpu_tensor = constant_op.constant(1.1)
result = cpu_tensor + cpu_tensor
self.assertEqual(result.device,
'/job:localhost/replica:0/task:0/device:GPU:0')
finally:
context._set_context(old_context)
def testSoftPlacement(self):
if not context.context().num_gpus():
self.skipTest('No GPUs found')
# Temporarily replace the context
# pylint: disable=protected-access
del context._context
context._context = context.Context()
try:
config.set_device_policy('silent')
config.set_soft_device_placement(True)
cpu_tensor = constant_op.constant(1.0)
result = cpu_tensor + cpu_tensor
self.assertEqual(result.device,
'/job:localhost/replica:0/task:0/device:GPU:0')
finally:
del context._context
context._context = context.Context()
def test_distribution(self, distribution):
input_array = np.array([[-1.5, 1.0, 3.4, .5], [0.0, 3.0, 1.3, 0.0]])
expected_output = [[0, 1, 3, 1], [0, 3, 2, 0]]
expected_output_shape = [None, 4]
config.set_soft_device_placement(True)
with distribution.scope():
input_data = keras.Input(shape=(4,))
layer = discretization.Discretization(bins=[0., 1., 2.])
bucket_data = layer(input_data)
self.assertAllEqual(expected_output_shape, bucket_data.shape.as_list())
model = keras.Model(inputs=input_data, outputs=bucket_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
def testVirtualGpu(self):
config.set_soft_device_placement(False)
gpus = config.list_physical_devices('GPU')
self.assertNotEqual(len(gpus), 0)
self.assertIsNone(config.get_logical_device_configuration(gpus[-1]))
config.set_logical_device_configuration(gpus[-1], [
context.LogicalDeviceConfiguration(memory_limit=10),
context.LogicalDeviceConfiguration(memory_limit=10)
])
self.assertEqual(
len(config.get_logical_device_configuration(gpus[-1])), 2)
logical_gpus = config.list_logical_devices('GPU')
self.assertTrue(len(logical_gpus), len(gpus) + 1)
for i in range(0, len(logical_gpus)):
with ops.device('/device:GPU:' + str(i)):
a = array_ops.identity(1.0)
self.evaluate(a)
with self.assertRaisesRegex(errors.InvalidArgumentError,
'Could not satisfy'):
with ops.device('/device:GPU:' + str(len(logical_gpus))):
a = array_ops.identity(1.0)
self.evaluate(a)
# Modifying the GPU configuration is not supported
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
config.set_logical_device_configuration(gpus[-1], [
context.LogicalDeviceConfiguration(memory_limit=20),
context.LogicalDeviceConfiguration(memory_limit=20)
])
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
config.set_logical_device_configuration(gpus[-1], [
context.LogicalDeviceConfiguration(memory_limit=10),
context.LogicalDeviceConfiguration(memory_limit=10),
context.LogicalDeviceConfiguration(memory_limit=10)
])
# Setting the same GPU configuration is fine
config.set_logical_device_configuration(gpus[-1], [
context.LogicalDeviceConfiguration(memory_limit=10),
context.LogicalDeviceConfiguration(memory_limit=10)
])
def test_distribution(self, distribution):
input_array_1 = np.array([['a', 'b'], ['c', 'd']])
input_array_2 = np.array([['e', 'f'], ['g', 'h']])
# pyformat: disable
expected_output = [[b'a_X_e', b'a_X_f', b'b_X_e', b'b_X_f'],
[b'c_X_g', b'c_X_h', b'd_X_g', b'd_X_h']]
config.set_soft_device_placement(True)
with distribution.scope():
input_data_1 = keras.Input(shape=(2, ), dtype=dtypes.string)
input_data_2 = keras.Input(shape=(2, ), dtype=dtypes.string)
input_data = [input_data_1, input_data_2]
layer = categorical_crossing.CategoryCrossing()
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict([input_array_1, input_array_2])
self.assertAllEqual(expected_output, output_dataset)
def test_distribution(self, distribution):
input_array = np.array([[1, 2, 3, 1], [0, 3, 1, 0]])
inp_dataset = dataset_ops.DatasetV2.from_tensor_slices(input_array)
inp_dataset = batch_wrapper(inp_dataset, 2, distribution)
# pyformat: disable
expected_output = [[0, 1, 1, 1, 0, 0], [1, 1, 0, 1, 0, 0]]
# pyformat: enable
max_tokens = 6
config.set_soft_device_placement(True)
with distribution.scope():
input_data = keras.Input(shape=(4, ), dtype=dtypes.int32)
layer = category_encoding.CategoryEncoding(
max_tokens=max_tokens, output_mode=category_encoding.BINARY)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(inp_dataset)
self.assertAllEqual(expected_output, output_dataset)
def test_dynamic_shape_with_outside_compilation_failure(self):
# Enable automatic outside compilation.
config.set_soft_device_placement(True)
strategy = get_tpu_strategy()
dataset = dataset_ops.Dataset.from_tensors(
("string", 1.0)).repeat().batch(2, drop_remainder=False)
dataset = strategy.experimental_distribute_dataset(dataset)
iterator = iter(dataset)
@def_function.function
def train_fn(iterator):
def step_fn(inputs):
_, inputs = inputs
return math_ops.reduce_sum(inputs)
return strategy.experimental_local_results(
strategy.run(step_fn, args=(next(iterator), )))
with self.assertRaisesRegex(errors.InternalError,
"Compilation failure"):
logging.info(train_fn(iterator))
def test_tpu_distribution(self, distribution):
vocab_data = [[
"earth", "earth", "earth", "earth", "wind", "wind", "wind", "and",
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(
input_array).batch(2, drop_remainder=True)
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
config.set_soft_device_placement(True)
with distribution.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = get_layer_class()()
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
def testGpuNone(self):
config.set_soft_device_placement(False)
gpus = config.list_physical_devices('GPU')
self.assertGreater(len(gpus), 0)
cpus = config.list_physical_devices('CPU')
self.assertEqual(len(cpus), 1)
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
self.assertGreater(len(config.get_visible_devices('GPU')), 0)
# get_visible_devices filters out XLA_* devices. list_logical_devices does
# not, but we can't call it here because it initializes the devices and
# calling set_visible_devices after that is disallowed.
self.assertEqual(len(config.get_visible_devices('XLA_GPU')), 0)
config.set_visible_devices(cpus[0])
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
self.assertEqual(len(config.get_visible_devices('GPU')), 0)
self.assertEqual(len(config.list_logical_devices('XLA_GPU')), 0)
with self.assertRaisesRegexp(errors.InvalidArgumentError,
'Could not satisfy'):
with ops.device('/device:GPU:0'):
a = array_ops.identity(1.0)
self.evaluate(a)
with self.assertRaisesRegexp(errors.InvalidArgumentError,
'Could not satisfy'):
with ops.device('/device:XLA_GPU:0'):
a = array_ops.identity(1.0)
self.evaluate(a)
# Modifying the visible devices is not supported
with self.assertRaisesRegexp(RuntimeError, 'cannot be modified'):
config.set_visible_devices(gpus)
# Setting the same visible devices is fine
config.set_visible_devices(cpus[0])
def test_single_tpu_jit_compile_with_outside_compilation(self):
context.enable_jit_compile_rewrite()
get_tpu_strategy(True)
config.set_soft_device_placement(True)
with ops.device("/device:TPU:0"):
a = variables.Variable(1)
def get_a_plus_one():
return a + 1
@def_function.function(
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
def foo(x):
b = x + get_a_plus_one()
my_str = string_ops.as_string(b)
new_str = my_str + "0"
c = string_ops.string_to_number(new_str, out_type=dtypes.int32)
logging_ops.print_v2(c)
b = c + get_a_plus_one()
return b + 1
with ops.device("/device:TPU:0"):
result = foo(a)
self.assertAllEqual(33, result)
def testEnableSoftPlacement(self):
self.assertEqual(config.get_soft_device_placement(), False)
config.set_soft_device_placement(True)
self.assertEqual(config.get_soft_device_placement(), True)
self.assertEqual(
config.get_soft_device_placement(),
context.context().soft_device_placement)
config.set_soft_device_placement(False)
self.assertEqual(config.get_soft_device_placement(), False)
self.assertEqual(
config.get_soft_device_placement(),
context.context().soft_device_placement)
constant_op.constant(1)
with self.assertRaises(RuntimeError):
config.set_soft_device_placement(True)
with self.assertRaises(RuntimeError):
config.set_soft_device_placement(False)
