# Copyright 2021 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. # ============================================================================== """Tests for the distributed variables library.""" import copy import os from absl.testing import parameterized from tensorflow.python.checkpoint import checkpoint as trackable_utils from tensorflow.python.checkpoint import checkpoint_options from tensorflow.python.distribute import collective_all_reduce_strategy from tensorflow.python.distribute import combinations from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import distribute_utils from tensorflow.python.distribute import packed_distributed_variable as packed from tensorflow.python.distribute import parameter_server_strategy from tensorflow.python.distribute import ps_values from tensorflow.python.distribute import strategy_combinations from tensorflow.python.distribute import test_util as ds_test_util from tensorflow.python.distribute import tpu_strategy from tensorflow.python.distribute import values as values_lib from tensorflow.python.eager import context from tensorflow.python.eager import def_function from tensorflow.python.eager import test from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import indexed_slices from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops from tensorflow.python.ops import control_flow_assert from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables as variables_lib from tensorflow.python.saved_model import save from tensorflow.python.saved_model import save_context from tensorflow.python.saved_model import save_options from tensorflow.python.types import core def _device_str(d): return "/device:GPU:" + str(d) def _nested_value(d): return ("a" + d, ["b" + d, {"c": "d" + d, "e": "f" + d}, "g" + d], "h" + d) def mirrored_and_tpu_strategy_combinations(): return combinations.combine( distribution=[ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call, strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, ], mode=["graph", "eager"]) def checkpoint_test_helper(dvar_test, distribution, synchronization, aggregation, enable_async_ckpt): # This method is added since `testCheckpointing` cannot be parameterized after # the entire class is parameterized. with distribution.scope(): v = variables_lib.Variable( constant_op.constant([1., 2., 3., 4]), synchronization=synchronization, aggregation=aggregation) dvar_test.evaluate(v.initializer) before_save = dvar_test.evaluate(v.read_value()) # Save random weights into checkpoint. checkpoint = trackable_utils.Checkpoint(v=v) ckpt_options = checkpoint_options.CheckpointOptions( experimental_enable_async_checkpoint=enable_async_ckpt) prefix = os.path.join(dvar_test.get_temp_dir(), "ckpt") with dvar_test.test_session(): save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options) # Assign inverted value. dvar_test.evaluate(v.assign(constant_op.constant([4., 3., 2., 1.]))) after_assign = dvar_test.evaluate(v.read_value()) dvar_test.assertNotAllClose(before_save, after_assign) # Restore from the checkpoint. with dvar_test.test_session(): checkpoint.restore(save_path).assert_consumed().run_restore_ops() after_restore = dvar_test.evaluate(v) dvar_test.assertAllClose(before_save, after_restore) # Another round of saving/restoring to ensure that the logic of # _copy_trackable_to_cpu works when a copy is already created in object_map. dvar_test.evaluate(v.assign(constant_op.constant([5., 6., 7., 8.]))) before_save_1 = dvar_test.evaluate(v.read_value()) dvar_test.assertNotAllClose(before_save_1, after_restore) with dvar_test.test_session(): save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options) dvar_test.evaluate(v.assign(constant_op.constant([8., 7., 6., 5.]))) after_assign_1 = dvar_test.evaluate(v.read_value()) dvar_test.assertNotAllClose(before_save_1, after_assign_1) with dvar_test.test_session(): checkpoint.restore(save_path).assert_consumed().run_restore_ops() after_restore_1 = dvar_test.evaluate(v) dvar_test.assertAllClose(before_save_1, after_restore_1) @combinations.generate( combinations.combine( distribution=[ strategy_combinations.mirrored_strategy_with_one_cpu, strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call, strategy_combinations.tpu_strategy, strategy_combinations.tpu_strategy_packed_var, strategy_combinations.tpu_strategy_spmd, strategy_combinations.central_storage_strategy_with_gpu_and_cpu, strategy_combinations.multi_worker_mirrored_2x1_cpu, strategy_combinations.multi_worker_mirrored_2x1_gpu, strategy_combinations.multi_worker_mirrored_2x2_gpu, strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call, ], synchronization=[ variables_lib.VariableSynchronization.ON_READ, variables_lib.VariableSynchronization.ON_WRITE, ], aggregation=[ variables_lib.VariableAggregation.MEAN, variables_lib.VariableAggregation.SUM, variables_lib.VariableAggregation.ONLY_FIRST_REPLICA, ], mode=["graph", "eager"], use_var_policy=[True, False])) class DistributedVariableTest(test.TestCase, parameterized.TestCase): def testExtendsVariable(self, distribution, synchronization, aggregation): with distribution.scope(): v = variables_lib.Variable( 1., synchronization=synchronization, aggregation=aggregation) self.assertIsInstance(v, variables_lib.Variable) def testCheckpointing(self, distribution, synchronization, aggregation, mode): if (isinstance(distribution, collective_all_reduce_strategy.CollectiveAllReduceStrategy) and mode == "graph"): self.skipTest("MWMS combinations tests do not work well in graph mode.") checkpoint_test_helper(self, distribution, synchronization, aggregation, enable_async_ckpt=False) def testAsyncCheckpointing(self, distribution, synchronization, aggregation, mode): if (isinstance(distribution, collective_all_reduce_strategy.CollectiveAllReduceStrategy) and mode == "graph"): self.skipTest("MWMS combinations tests do not work well in graph mode.") checkpoint_test_helper(self, distribution, synchronization, aggregation, enable_async_ckpt=True) def testTraceback(self, distribution, synchronization, aggregation): if context.executing_eagerly(): self.skipTest("does not apply to eager") with distribution.scope(): variable_scope.get_variable( name="testVar", initializer=1., use_resource=True, synchronization=synchronization, aggregation=aggregation) with self.assertRaisesRegex(ValueError, "Variable testVar already exists"): variable_scope.get_variable( name="testVar", initializer=1., use_resource=True, synchronization=synchronization, aggregation=aggregation) def testSelectReplica(self, distribution, synchronization, aggregation): with distribution.scope(): v = variables_lib.Variable( 1., synchronization=synchronization, aggregation=aggregation) self.assertIs(v, distribute_utils.select_replica(0, v)) def testIsTensorLike(self, distribution, synchronization, aggregation): if isinstance(distribution.extended, tpu_strategy.TPUExtended) and context.executing_eagerly(): self.skipTest("TPU doesn't support pure eager") with distribution.scope(): v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation) # In cross replica context. self.assertIsInstance(v, core.Tensor) # In replica context. distribution.run(lambda v: self.assertIsInstance(v, core.Tensor), args=(v,)) def testAssignReturnValueIsTensorLike(self, distribution, synchronization, aggregation): if isinstance(distribution.extended, tpu_strategy.TPUExtended): if context.executing_eagerly(): self.skipTest("TPU doesn't support pure eager") else: self.skipTest("b/152076846") with distribution.scope(): v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation) def assert_is_tensor_like(v): # We can't use Python literals because they are treated as non-distributed # values is not allowed when aggregation is SUM. See # `cross_device_ops.reduce_non_distributed_value`. delta = array_ops.identity(1.) self.assertIsInstance(v.assign(delta), core.Tensor) self.assertIsInstance(v.assign_sub(delta), core.Tensor) self.assertIsInstance(v.assign_add(delta), core.Tensor) # In cross replica context we return a PerReplica which is not Tensor like # all the time yet. if (synchronization == variables_lib.VariableSynchronization.ON_READ and aggregation != variables_lib.VariableAggregation.SUM): assert_is_tensor_like(v) # In replica context. distribution.run(assert_is_tensor_like, args=(v,)) def testDeepCopy(self, distribution, synchronization, aggregation): if not context.executing_eagerly(): self.skipTest("deepcopy only supported in eager mode") with distribution.scope(): v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation) in_dist_copy = copy.deepcopy(v) out_dist_copy = copy.deepcopy(v) def assert_is_deep_copy(v1, v2): self.assertIsInstance(v2, type(v1)) self.assertEqual(v1.aggregation, v2.aggregation) self.assertEqual(v1.distribute_strategy, v2.distribute_strategy) if isinstance(v1, ps_values.AggregatingVariable): self.assertIsInstance(v2.get(), type(v1.get())) self.assertNotEqual(id(v1.get()), id(v2.get())) else: if v1._policy: self.assertNotEqual(id(v1._policy), id(v2._policy)) # pylint: disable=protected-access else: self.assertEqual(id(v1._policy), id(v2._policy)) # pylint: disable=protected-access self.assertEqual(len(v1.values), len(v2.values)) for (v1v, v2v) in zip(v1.values, v2.values): self.assertEqual(v1v.device, v2v.device) self.assertNotEqual(id(v1v), id(v2v)) self.assertAllEqual( self.evaluate(v1.values), self.evaluate(v2.values)) self.evaluate(variables_lib.global_variables_initializer()) if not isinstance(distribution.extended, tpu_strategy.TPUExtended): distribution.run(assert_is_deep_copy, args=(v, in_dist_copy)) distribution.run(assert_is_deep_copy, args=(v, out_dist_copy)) def testAssignSignature(self, distribution, synchronization, aggregation): # This test verifies assign*() can be called in the same way as normal # variables. with distribution.scope(): v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation) def assign(): one = constant_op.constant(1.) v.assign(one, True, "assign", False) # TODO(b/154017756): SyncOnReadVariable.assign() doesn't support passing # value as a keyword argument. v.assign(one, use_locking=True, name="assign", read_value=False) v.assign_add(one, True, "assign", False) v.assign_add(one, use_locking=True, name="assign", read_value=False) v.assign_sub(one, True, "assign", False) v.assign_sub(one, use_locking=True, name="assign", read_value=False) # Return something for graph mode to fetch. return constant_op.constant(1) self.evaluate(variables_lib.global_variables_initializer()) if not (synchronization == variables_lib.VariableSynchronization.ON_READ and aggregation == variables_lib.VariableAggregation.SUM): self.evaluate(distribution.experimental_local_results(assign())) if not (isinstance(distribution.extended, tpu_strategy.TPUExtended) and context.executing_eagerly()): self.evaluate( distribution.experimental_local_results(distribution.run(assign))) def testStrategyExtendedUpdate(self, distribution, synchronization, aggregation): if len(distribution.extended.parameter_devices) != 2: self.skipTest("n/a: needs exactly two parameter devices") if (synchronization == variables_lib.VariableSynchronization.ON_WRITE and aggregation != variables_lib.VariableAggregation.NONE): self.skipTest("n/a: doesn't apply to ON_WRITE variable with aggregation") with distribution.scope(): v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation) value = values_lib.PerReplica([1., 2.]) assign_fn = lambda var, value: var.assign(value) self.evaluate(distribution.extended.update(v, assign_fn, args=(value,))) self.assertAllEqual(self.evaluate(v.values), [1., 2.]) assign_add_fn = lambda var, value: var.assign_add(value) self.evaluate(distribution.extended.update(v, assign_add_fn, args=(value,))) self.assertAllEqual(self.evaluate(v.values), [2., 4.]) assign_sub_fn = lambda var, value: var.assign_sub(value) self.evaluate(distribution.extended.update(v, assign_sub_fn, args=(value,))) self.assertAllEqual(self.evaluate(v.values), [1., 2.]) read_assign_fn = lambda var, value: var.assign_add(var.value() + var. read_value()) self.evaluate( distribution.extended.update(v, read_assign_fn, args=(value,))) self.assertAllEqual(self.evaluate(v.values), [3., 6.]) def testSaveNonDistributed(self, distribution, synchronization, aggregation): # This test verifies that the DistributedVariable behave like the primary # variable when saving a non-distributed version of the model (the default). # The test asserts that the function traced under SaveContext has no device # annotations and only reference the primary component of the variable. Note # that please avoid capturing other eager tensors in this test to make the # assertion easy. if isinstance(distribution.extended, parameter_server_strategy.ParameterServerStrategyExtended): self.skipTest("b/148689177: AggregatingVariable doesn't " "conform to Variable interface well") # tf.function requires the return value to be Tensors, which is not always # case for properties and methods of Variable, so we simply discard the # return values. def _discard_return(f): f() return def _test(f, v): # This verifies that the function under SaveContext: # - contains no device annotations. # - only references the primary component of the variable. g = def_function.function(lambda: _discard_return(f)) options = save_options.SaveOptions( experimental_variable_policy=save_options.VariablePolicy.NONE) with save_context.save_context(options): # The graph should contain no device. graph = g.get_concrete_function().graph for op in graph.get_operations(): self.assertEqual(op.device, "", msg=str(op)) # The function should only capture the primary variable. Note that it # may not have captures, e.g. v.aggregation. captures = list(graph.captures) self.assertLessEqual(len(captures), 1) if graph.captures: self.assertIs(captures[0][0], v._primary.handle) def _assert(cond): return control_flow_assert.Assert(cond, [cond]) with distribution.scope(): # We use four variables for convenience reasons. They have no special # meaning. # - v is used whenever possible. # - w is used for scatter and gather, which require the variable to be # non-scalar. # - y is used when the dtype needs to be integer. Note that aggregation # cannot be MEAN for integers. v = variables_lib.Variable( 0., synchronization=synchronization, aggregation=aggregation, trainable=True) w = variables_lib.Variable([0., 0., 0.], synchronization=synchronization, aggregation=aggregation, trainable=True) if aggregation != variables_lib.VariableAggregation.MEAN: y = variables_lib.Variable( 0, synchronization=synchronization, aggregation=aggregation) # pylint: disable=g-long-lambda # tf.Variable properties. _test(lambda: self.assertEqual(v.aggregation, aggregation), v) _test(lambda: self.assertIs(v.constraint, None), v) # TODO(crccw): should we raise an error instead? _test(lambda: self.assertEqual(v.device, v._primary.device), v) _test(lambda: self.assertEqual(v.dtype, dtypes.float32), v) if not context.executing_eagerly(): _test(lambda: self.assertIs(v.graph, v._primary.graph), v) if not context.executing_eagerly(): _test(lambda: _assert(v.initial_value == 0), v) _test(lambda: self.assertIs(v.initializer, v._primary.initializer), v) _test(lambda: self.assertEqual(v.name, "Variable:0"), v) if not context.executing_eagerly(): _test(lambda: self.assertIs(v.op, v._primary.op), v) _test(lambda: self.assertEqual(v.shape, tensor_shape.TensorShape(())), v) _test(lambda: self.assertEqual(v.synchronization, synchronization), v) _test(lambda: self.assertEqual(v.trainable, True), v) # tf.Variable methods. _test(lambda: check_ops.assert_equal_v2(v.assign(1.), 1.), v) _test(lambda: check_ops.assert_equal_v2(v.assign_add(1.), 2.), v) _test(lambda: check_ops.assert_equal_v2(v.assign_sub(1.), 1.), v) # TODO(b/148689177): Implement batch_scatter_update. # count_up_to() is skipped since it's deprecated. # eval() is skipped since it shouldn't called in a tf.function. # experimental_ref() is skipped since it's deprecated. # from_proto() is skipped since it shouldn't called in a tf.function. # TODO(b/148689177): Implement gather_nd. _test( lambda: check_ops.assert_equal_v2(v.get_shape(), tensor_shape.TensorShape(())), v) # initialized_value() is skipped since it shouldn't called in a tf.function. # load() is skipped since it shouldn't called in a tf.function. _test(lambda: check_ops.assert_equal_v2(v.read_value(), 1.), v) # ref() is skipped since it shouldn't called in a tf.function. _test( lambda: check_ops.assert_equal_v2( w.scatter_add(_make_index_slices(values=[1., 2.], indices=[0, 2])), [1., 0., 2.]), w) _test( lambda: check_ops.assert_equal_v2( w.scatter_div(_make_index_slices(values=[4., 2.], indices=[0, 2])), [0.25, 0., 1.]), w) _test( lambda: check_ops.assert_equal_v2( w.scatter_max(_make_index_slices(values=[1., 0.5], indices=[1, 2])), [0.25, 1., 1.]), w) _test( lambda: check_ops.assert_equal_v2( w.scatter_min(_make_index_slices(values=[1., 0.5], indices=[0, 1])), [0.25, 0.5, 1.]), w) _test( lambda: check_ops.assert_equal_v2( w.scatter_mul(_make_index_slices(values=[2., 0.5], indices=[0, 1])), [0.5, 0.25, 1.]), w) # TODO(b/148689177): Implement scatter_nd_* _test( lambda: check_ops.assert_equal_v2( w.scatter_sub(_make_index_slices(values=[2., 0.5], indices=[0, 1])), [-1.5, -0.25, 1.]), w) _test( lambda: check_ops.assert_equal_v2( w.scatter_update( _make_index_slices(values=[2., 0.5], indices=[0, 1])), [2., 0.5, 1.]), w) # set_shape() is skipped since ResourceVariable doesn't implement it. # to_proto() is skipped since it shouldn't called in a tf.function. _test(lambda: check_ops.assert_equal_v2(v.value(), 1.), v) # DistributedVariable should be treated as ResourceVariable, so it needs to # conform to ResourceVariable interface as well. _test(lambda: self.assertIs(v.handle, v._primary.handle), v) # Convert to tensor. _test(lambda: check_ops.assert_equal_v2(ops.convert_to_tensor(v), 1.), v) # Control dependency. def _with_control_dep(): with ops.control_dependencies([v.assign(1.)]): return array_ops.identity(1) _test(_with_control_dep, v) # Operator overloads. _test(lambda: check_ops.assert_equal_v2(v.assign(7.), 7.), v) _test(lambda: check_ops.assert_equal_v2(v + 1., 8.), v) _test(lambda: check_ops.assert_equal_v2(3 + v, 10.), v) _test(lambda: check_ops.assert_equal_v2(v + v, 14.), v) _test(lambda: check_ops.assert_equal_v2(v - 2., 5.), v) _test(lambda: check_ops.assert_equal_v2(v - v, 0.), v) _test(lambda: check_ops.assert_equal_v2(v * 2., 14.), v) _test(lambda: check_ops.assert_equal_v2(3 * v, 21.), v) _test(lambda: check_ops.assert_equal_v2(v * v, 49.), v) _test( lambda: check_ops.assert_equal_v2( math_ops.cast(v / 2., dtypes.float32), 3.5), v) _test( lambda: check_ops.assert_equal_v2( math_ops.cast(14. / v, dtypes.float32), 2.), v) _test(lambda: _assert(v < 12.), v) _test(lambda: _assert(v <= 12.), v) _test(lambda: _assert(not v > 12.), v) _test(lambda: _assert(not v >= 12.), v) _test(lambda: _assert(not 12. < v), v) _test(lambda: _assert(not 12. <= v), v) _test(lambda: _assert(12. > v), v) _test(lambda: _assert(12. >= v), v) _test(lambda: check_ops.assert_near_v2(pow(v, 3.), 343.), v) _test(lambda: check_ops.assert_near_v2(pow(2., v), 128.), v) _test(lambda: check_ops.assert_equal_v2(abs(v), 7.), v) # Operator overloads that only works for integers. if aggregation != variables_lib.VariableAggregation.MEAN: _test(lambda: check_ops.assert_equal_v2(y.assign(7), 7), y) _test(lambda: check_ops.assert_equal_v2(y // 2, 3), y) _test(lambda: check_ops.assert_equal_v2(15 // y, 2), y) _test(lambda: check_ops.assert_equal_v2(y % 2, 1), y) _test(lambda: check_ops.assert_equal_v2(16 % y, 2), y) _test(lambda: check_ops.assert_equal_v2(y & 3, 3), y) _test(lambda: check_ops.assert_equal_v2(3 & y, 3), y) _test(lambda: check_ops.assert_equal_v2(y | 8, 15), y) _test(lambda: check_ops.assert_equal_v2(16 | y, 23), y) _test(lambda: check_ops.assert_equal_v2(y ^ 3, 4), y) _test(lambda: check_ops.assert_equal_v2(11 ^ y, 12), y) _test(lambda: check_ops.assert_equal_v2(-y, -7), y) _test(lambda: check_ops.assert_equal_v2(~y, ~7), y) # Index. if isinstance(distribution.extended, tpu_strategy.TPUExtended): # TODO(b/161572567): slice assignment doesn't work for TPU. _test(lambda: check_ops.assert_equal_v2(w[0], 2.), w) else: _test(lambda: check_ops.assert_equal_v2(w[0].assign(1.), [1., 0.5, 1.]), w) _test(lambda: check_ops.assert_equal_v2(w[0], 1.), w) # pylint: enable=g-long-lambda def testUnsaveable(self, distribution, synchronization, aggregation, mode): if isinstance(distribution.extended, parameter_server_strategy.ParameterServerStrategyExtended): self.skipTest("n/a: not applicable to AggregatingVariable") if (isinstance(distribution, collective_all_reduce_strategy.CollectiveAllReduceStrategy) and mode == "graph"): self.skipTest("MWMS combinations tests do not work well in graph mode.") if not distribution.extended._use_merge_call(): self.skipTest("Unsupported combination.") with distribution.scope(): v = variables_lib.Variable([1., 1.], synchronization=synchronization, aggregation=aggregation) with self.cached_session(): self.evaluate(variables_lib.global_variables_initializer()) export_dir = self.get_temp_dir() def _assert_unsaveable(f): # Ignore if it cannot be traced. Certain combinations are not supported or # yet or not allowed. try: f = def_function.function(f).get_concrete_function() except (NotImplementedError, ValueError): return with self.assertRaisesRegex(ValueError, "f_with_input_signature"): save.save(v, export_dir, signatures=f) _assert_unsaveable(lambda: v.assign(ops.convert_to_tensor([1., 1.]))) _assert_unsaveable(lambda: v.assign_add(ops.convert_to_tensor([1., 1.]))) _assert_unsaveable(lambda: v.assign_sub(ops.convert_to_tensor([1., 1.]))) _assert_unsaveable(lambda: v.scatter_add(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_sub(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_mul(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_div(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_min(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_max(_make_index_slices([1.], [0]))) _assert_unsaveable(lambda: v.scatter_update(_make_index_slices([1.], [0]))) # Reading a ON_READ variable should be unsaveable if either: # 1) CollectiveAllReduceStrategy, and aggregation is MEAN/SUM. # 2) aggregation is SUM. if (synchronization == variables_lib.VariableSynchronization.ON_READ and (aggregation == variables_lib.VariableAggregation.SUM or (not distribution.extended._use_merge_call()) or (isinstance(distribution.extended, collective_all_reduce_strategy.CollectiveAllReduceExtended) and aggregation == variables_lib.VariableAggregation.MEAN))): _assert_unsaveable(v.read_value) _assert_unsaveable(v.value) _assert_unsaveable(lambda: ops.convert_to_tensor(v)) else: # Otherwise reading a variable should be saveable. @def_function.function def f(): v.read_value() v.value() return ops.convert_to_tensor(v) with self.cached_session(): save.save(v, export_dir, signatures=f.get_concrete_function()) @combinations.generate( combinations.combine( distribution=[ strategy_combinations.mirrored_strategy_with_one_cpu, strategy_combinations.tpu_strategy, ], mode=["eager"])) class PackedDistributedVariableTest(test.TestCase, parameterized.TestCase): def testPackedVariable(self, distribution): with distribution.scope(): v0 = variables_lib.Variable(0.) self.assertIsNone(v0._packed_var) distribution._enable_packed_variable_in_eager_mode = True with distribution.scope(): v1 = variables_lib.Variable(0) self.assertIsInstance(v1._packed_var, packed.PackedDistributedVariable) devices = v1._devices for i in range(1, len(devices)): with distribute_lib.ReplicaContext(distribution, i): v1.assign(i) val = v1._get() self.assertIsInstance(val, packed.PackedVarAndDevice) self.assertEqual(val.device, devices[0]) self.assertEqual(self.evaluate(val.read_value()), 0) for i in range(0, len(devices)): with distribute_lib.ReplicaContext(distribution, i): val = v1._get() self.assertIsInstance(val, packed.PackedVarAndDevice) self.assertEqual(val.device, devices[i]) self.assertEqual(self.evaluate(val.read_value()), i) def testIgnorePackedVariableInSaveContext(self, distribution): distribution._enable_packed_variable_in_eager_mode = True with distribution.scope(): v = variables_lib.Variable(0) self.assertIsInstance(v._packed_variable, packed.PackedDistributedVariable) options = save_options.SaveOptions() with save_context.save_context(options): self.assertIsNone(v._packed_variable) def _make_index_slices(values, indices, dense_shape=None): if dense_shape: dense_shape = array_ops.identity(dense_shape) return indexed_slices.IndexedSlices( array_ops.identity(values), array_ops.identity(indices), dense_shape) if __name__ == "__main__": ds_test_util.main()