chore: import upstream snapshot with attribution

tensorflow/python/kernel_tests/variables/BUILD

@@ -0,0 +1,200 @@
+# Tests of TensorFlow variables kernels written using the Python API.
+
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cuda_py_strict_test(
+    name = "dense_update_ops_no_tsan_test",
+    size = "small",
+    srcs = ["dense_update_ops_no_tsan_test.py"],
+    # TODO (b/140294007): the test fails with XLA.
+    xla_enable_strict_auto_jit = False,
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "dense_update_ops_test",
+    size = "small",
+    srcs = ["dense_update_ops_test.py"],
+    deps = [
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "partitioned_variables_test",
+    size = "small",
+    srcs = ["partitioned_variables_test.py"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:partitioned_variables",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:saver",
+        "//third_party/py/numpy",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "resource_variable_ops_test",
+    size = "medium",
+    srcs = ["resource_variable_ops_test.py"],
+    # TODO(b/128347673): Re-enable.
+    tags = [
+        "no_windows",
+        "nozapfhahn",
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compat",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:composite_tensor_gradient",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:cpp_shape_inference_proto_py",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:extension_type",
+        "//tensorflow/python/framework:indexed_slices",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/framework:test_ops",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:custom_gradient",
+        "//tensorflow/python/ops:gradients_impl",
+        "//tensorflow/python/ops:handle_data_util",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:momentum",
+        "//tensorflow/python/training:saver",
+        "//tensorflow/python/training:training_util",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "variable_ops_test",
+    size = "small",
+    srcs = ["variable_ops_test.py"],
+    tags = ["nomsan"],  # TODO(b/281849652): Re-enable.
+    deps = [
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:state_ops_gen",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "variable_scope_test",
+    size = "medium",
+    srcs = ["variable_scope_test.py"],
+    tags = ["no_windows"],
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:wrap_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:init_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:resource_variables_toggle",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:tf_inspect",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "variables_test",
+    size = "small",
+    srcs = ["variables_test.py"],
+    tags = ["no_windows"],  # b/133869052
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:state_ops_gen",
+        "//tensorflow/python/ops:variable_v1",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:gradient_descent",
+        "//tensorflow/python/util:compat",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/kernel_tests/variables/dense_update_ops_no_tsan_test.py

@@ -0,0 +1,163 @@
+# Copyright 2015 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 state updating ops that may have benign race conditions."""
+
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class AssignOpTest(test.TestCase):
+
+  # NOTE(mrry): We exclude thess tests from the TSAN TAP target, because they
+  #   contain benign and deliberate data races when multiple threads update
+  #   the same parameters without a lock.
+  def testParallelUpdateWithoutLocking(self):
+    # We need each thread to keep its own device stack or the device scopes
+    # won't be properly nested.
+    ops.get_default_graph().switch_to_thread_local()
+    with self.cached_session() as sess:
+      ones_t = array_ops.fill([1024, 1024], 1.0)
+      p = variables.Variable(array_ops.zeros([1024, 1024]))
+      adds = [
+          state_ops.assign_add(
+              p, ones_t, use_locking=False) for _ in range(20)
+      ]
+      self.evaluate(variables.global_variables_initializer())
+
+      def run_add(add_op):
+        self.evaluate(add_op)
+
+      threads = [
+          self.checkedThread(
+              target=run_add, args=(add_op,)) for add_op in adds
+      ]
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+
+      vals = self.evaluate(p)
+      ones = np.ones((1024, 1024)).astype(np.float32)
+      self.assertTrue((vals >= ones).all())
+      self.assertTrue((vals <= ones * 20).all())
+
+  def testParallelAssignWithoutLocking(self):
+    # We need each thread to keep its own device stack or the device scopes
+    # won't be properly nested.
+    ops.get_default_graph().switch_to_thread_local()
+    with self.cached_session() as sess:
+      ones_t = array_ops.fill([1024, 1024], float(1))
+      p = variables.Variable(array_ops.zeros([1024, 1024]))
+      assigns = [
+          state_ops.assign(p, math_ops.multiply(ones_t, float(i)), False)
+          for i in range(1, 21)
+      ]
+      self.evaluate(variables.global_variables_initializer())
+
+      def run_assign(assign_op):
+        self.evaluate(assign_op)
+
+      threads = [
+          self.checkedThread(
+              target=run_assign, args=(assign_op,)) for assign_op in assigns
+      ]
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+
+      vals = self.evaluate(p)
+
+      # Assert every element is taken from one of the assignments.
+      self.assertTrue((vals > 0).all())
+      self.assertTrue((vals <= 20).all())
+
+  # NOTE(skyewm): We exclude these tests from the TSAN TAP target, because they
+  # contain non-benign but known data races between the variable assignment and
+  # returning the output tensors. This issue will be resolved with the new
+  # resource variables.
+  def testParallelUpdateWithLocking(self):
+    # We need each thread to keep its own device stack or the device scopes
+    # won't be properly nested.
+    ops.get_default_graph().switch_to_thread_local()
+    with self.cached_session() as sess:
+      zeros_t = array_ops.fill([1024, 1024], 0.0)
+      ones_t = array_ops.fill([1024, 1024], 1.0)
+      p = variables.Variable(zeros_t)
+      adds = [
+          state_ops.assign_add(
+              p, ones_t, use_locking=True) for _ in range(20)
+      ]
+      self.evaluate(p.initializer)
+
+      def run_add(add_op):
+        self.evaluate(add_op)
+
+      threads = [
+          self.checkedThread(
+              target=run_add, args=(add_op,)) for add_op in adds
+      ]
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+
+      vals = self.evaluate(p)
+      ones = np.ones((1024, 1024)).astype(np.float32)
+      self.assertAllEqual(vals, ones * 20)
+
+  def testParallelAssignWithLocking(self):
+    # We need each thread to keep its own device stack or the device scopes
+    # won't be properly nested.
+    ops.get_default_graph().switch_to_thread_local()
+    with self.cached_session() as sess:
+      zeros_t = array_ops.fill([1024, 1024], 0.0)
+      ones_t = array_ops.fill([1024, 1024], 1.0)
+      p = variables.Variable(zeros_t)
+      assigns = [
+          state_ops.assign(
+              p, math_ops.multiply(ones_t, float(i)), use_locking=True)
+          for i in range(1, 21)
+      ]
+      self.evaluate(p.initializer)
+
+      def run_assign(assign_op):
+        self.evaluate(assign_op)
+
+      threads = [
+          self.checkedThread(
+              target=run_assign, args=(assign_op,)) for assign_op in assigns
+      ]
+      for t in threads:
+        t.start()
+      for t in threads:
+        t.join()
+
+      vals = self.evaluate(p)
+
+      # Assert every element is the same, and taken from one of the assignments.
+      self.assertTrue(vals[0, 0] > 0)
+      self.assertTrue(vals[0, 0] <= 20)
+      self.assertAllEqual(vals, np.ones([1024, 1024]) * vals[0, 0])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/variables/dense_update_ops_test.py

@@ -0,0 +1,87 @@
+# Copyright 2015 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 tensorflow.ops.tf.Assign*."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class AssignOpTest(test.TestCase):
+
+  def _initAssignFetch(self, x, y, use_gpu):
+    """Initialize a param to init and update it with y."""
+    super(AssignOpTest, self).setUp()
+    with test_util.device(use_gpu=use_gpu):
+      p = variables.Variable(x)
+      assign = state_ops.assign(p, y)
+      self.evaluate(p.initializer)
+      new_value = self.evaluate(assign)
+      return self.evaluate(p), new_value
+
+  def _initAssignAddFetch(self, x, y, use_gpu):
+    """Initialize a param to init, and compute param += y."""
+    with test_util.device(use_gpu=use_gpu):
+      p = variables.Variable(x)
+      add = state_ops.assign_add(p, y)
+      self.evaluate(p.initializer)
+      new_value = self.evaluate(add)
+      return self.evaluate(p), new_value
+
+  def _initAssignSubFetch(self, x, y, use_gpu):
+    """Initialize a param to init, and compute param -= y."""
+    with test_util.device(use_gpu=use_gpu):
+      p = variables.Variable(x)
+      sub = state_ops.assign_sub(p, y)
+      self.evaluate(p.initializer)
+      new_value = self.evaluate(sub)
+      return self.evaluate(p), new_value
+
+  def _testTypes(self, vals):
+    for dtype in [
+        np.float32, np.float64, np.int32, np.int64,
+        dtypes.bfloat16.as_numpy_dtype
+    ]:
+      x = np.zeros(vals.shape).astype(dtype)
+      y = vals.astype(dtype)
+      var_value, op_value = self._initAssignFetch(x, y, use_gpu=False)
+      self.assertAllEqual(y, var_value)
+      self.assertAllEqual(y, op_value)
+      var_value, op_value = self._initAssignAddFetch(x, y, use_gpu=False)
+      self.assertAllEqual(x + y, var_value)
+      self.assertAllEqual(x + y, op_value)
+      var_value, op_value = self._initAssignSubFetch(x, y, use_gpu=False)
+      self.assertAllEqual(x - y, var_value)
+      self.assertAllEqual(x - y, op_value)
+      if test.is_built_with_gpu_support() and dtype in [np.float32, np.float64]:
+        var_value, op_value = self._initAssignFetch(x, y, use_gpu=True)
+        self.assertAllEqual(y, var_value)
+        self.assertAllEqual(y, op_value)
+        var_value, op_value = self._initAssignAddFetch(x, y, use_gpu=True)
+        self.assertAllEqual(x + y, var_value)
+        self.assertAllEqual(x + y, op_value)
+        var_value, op_value = self._initAssignSubFetch(x, y, use_gpu=True)
+        self.assertAllEqual(x - y, var_value)
+        self.assertAllEqual(x - y, op_value)
+
+  def testBasic(self):
+    self._testTypes(np.arange(0, 20).reshape([4, 5]))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/variables/partitioned_variables_test.py

@@ -0,0 +1,642 @@
+# Copyright 2015 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 partitioned_variables.py."""
+
+import os
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import saver as saver_lib
+
+
+def initialized_value(var):
+  return cond.cond(
+      variable_v1.is_variable_initialized(var), var.read_value,
+      lambda: var.initial_value)
+
+
+class PartitionerCreatorsTest(test.TestCase):
+
+  def testFixedSizePartitioner(self):
+    with self.cached_session():
+      partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
+      with variable_scope.variable_scope("root", partitioner=partitioner):
+        v0 = variable_scope.get_variable(
+            "v0", dtype=dtypes.float32, shape=(10, 10))
+        v0_list = v0._get_variable_list()
+        v0_part = v0._get_partitions()
+        self.assertEqual(len(v0_list), 5)
+        self.assertAllEqual(v0_part, (5, 1))
+
+  def testFixedSizePartitionerInt64(self):
+    with self.cached_session():
+      partitioner = partitioned_variables.fixed_size_partitioner(4, axis=0)
+      with variable_scope.variable_scope("root", partitioner=partitioner):
+        v0 = variable_scope.get_variable("v0", dtype=dtypes.int64, shape=[20])
+        v0_list = v0._get_variable_list()
+        self.assertEqual(len(v0_list), 4)
+
+  def testResourceFixedSizePartitioner(self):
+    with self.cached_session():
+      partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
+      with variable_scope.variable_scope(
+          "root", partitioner=partitioner, use_resource=True):
+        v0 = variable_scope.get_variable(
+            "v0", dtype=dtypes.float32, shape=(10, 10))
+        v0_list = v0._get_variable_list()
+        v0_part = v0._get_partitions()
+        self.assertEqual(len(v0_list), 5)
+        self.assertAllEqual(v0_part, (5, 1))
+
+  def _testVariableAxisSizePartitioner(self,
+                                       name,
+                                       axis,
+                                       max_shard_bytes,
+                                       expected_axis_shards,
+                                       expected_partitions,
+                                       max_shards=None):
+    partitioner = partitioned_variables.variable_axis_size_partitioner(
+        axis=axis, max_shard_bytes=max_shard_bytes, max_shards=max_shards)
+
+    with variable_scope.variable_scope("root", partitioner=partitioner):
+      v0 = variable_scope.get_variable(
+          name, dtype=dtypes.float32, shape=(4, 8, 16, 32))
+      v0_list = v0._get_variable_list()
+      v0_part = v0._get_partitions()
+      self.assertEqual(len(v0_list), expected_axis_shards)
+      self.assertAllEqual(v0_part, expected_partitions)
+
+  def testVariableAxisSizePartitioner(self):
+    with self.cached_session():
+      # Create a partitioned variable of shape (4, 8, 16, 32) type float32
+      # Bytes per slice along the given axes:
+
+      # 8 * 16 * 32 * sizeof(float32) = 16384 / slice on axis 0
+      # 4 * 16 * 32 * sizeof(float32) = 8192 / slice on axis 1
+      # 4 * 8 * 32 * sizeof(float32) = 4096 / slice on axis 2
+      # 4 * 8 * 16 * sizeof(float32) = 2048 / slice on axis 3
+
+      # Now partition it in different ways...
+
+      # No need to slice: bytes_per_slice * dim0 = 65536 < max_shard_bytes
+      self._testVariableAxisSizePartitioner(
+          "v0",
+          axis=0,
+          max_shard_bytes=131072,
+          expected_axis_shards=1,
+          expected_partitions=(1, 1, 1, 1))
+
+      # Slice exactly once: bytes_per_slice * dim1 = 65536 = max_shard_bytes
+      self._testVariableAxisSizePartitioner(
+          "v1",
+          axis=1,
+          max_shard_bytes=65536,
+          expected_axis_shards=1,
+          expected_partitions=(1, 1, 1, 1))
+
+      # Slice into 2 parts:
+      # bytes_per_slice = 4096
+      # slices_per_shard = 32768 / 4096 = 8
+      # axis_shards = 16 / 8 = 2
+      self._testVariableAxisSizePartitioner(
+          "v2",
+          axis=2,
+          max_shard_bytes=32768,
+          expected_axis_shards=2,
+          expected_partitions=(1, 1, 2, 1))
+
+      # This partitioner makes sure we maximize the number of shards along
+      # axis 3. Slice it into 32 parts:
+      # bytes_per_slice = 2048
+      # slices_per_shard = 2048 / 2048 = 1
+      # axis_shards = 32 / 1 = 32
+      self._testVariableAxisSizePartitioner(
+          "v3a",
+          axis=3,
+          max_shard_bytes=2048,
+          expected_axis_shards=32,
+          expected_partitions=(1, 1, 1, 32))
+
+      # This partitioner makes sure we do not go past the bound of allowable
+      # number of shards along axis 3.
+      # Slice into 32 parts:
+      # bytes_per_slice = 2048
+      # slices_per_shard = max(1, 1024 / 2048) = 1
+      # axis_shards = 32 / 1 = 32
+      # Slice into max of 32 parts because: max_shard_bytes < bytes_per_slice
+      self._testVariableAxisSizePartitioner(
+          "v3b",
+          axis=3,
+          max_shard_bytes=1024,
+          expected_axis_shards=32,
+          expected_partitions=(1, 1, 1, 32))
+
+      # Specify max_shards so that it won't affect sharding.
+      self._testVariableAxisSizePartitioner(
+          "v3c",
+          axis=3,
+          max_shard_bytes=1024,
+          expected_axis_shards=32,
+          expected_partitions=(1, 1, 1, 32),
+          max_shards=33)
+
+      # Specify max_shards so that it will affect sharding.
+      self._testVariableAxisSizePartitioner(
+          "v3d",
+          axis=3,
+          max_shard_bytes=1024,
+          expected_axis_shards=2,
+          expected_partitions=(1, 1, 1, 2),
+          max_shards=2)
+
+      # Use the partitioner with strings
+      partitioner_axis3_str = partitioned_variables.variable_axis_size_partitioner(  # pylint: disable=line-too-long
+          axis=3,
+          max_shard_bytes=32768,
+          bytes_per_string_element=8)
+
+      with variable_scope.variable_scope(
+          "root", partitioner=partitioner_axis3_str):
+        v3str = variable_scope.get_variable(
+            "v3str",
+            initializer=np.array([""] * 4 * 8 * 16 * 32).reshape(4, 8, 16, 32),  # pylint: disable=too-many-function-args
+            dtype=dtypes.string,
+            shape=(4, 8, 16, 32))
+        v3str_list = v3str._get_variable_list()
+        v3str_part = v3str._get_partitions()
+
+        # Now the estimated bytes_per_slice = 4*8*16*bytes_per_string_element
+        # which is equal to 4096.  Setting a max_shard_bytes of 32768
+        # and we should get a split of 4.
+        # Slice into 4 parts:
+        # bytes_per_slice = 4096
+        # slices_per_shard = 32768 / 4096 = 8
+        # axis_shards = 32 / 8 = 4
+        self.assertEqual(len(v3str_list), 4)
+        self.assertAllEqual(v3str_part, (1, 1, 1, 4))
+
+  def _testMinMaxVariablePartitioner(self, max_partitions, axis, min_slice_size,
+                                     var_name, var_shape, expected_axis_shards,
+                                     expected_partitions):
+    partitioner = partitioned_variables.min_max_variable_partitioner(
+        max_partitions=max_partitions, axis=axis, min_slice_size=min_slice_size)
+    with variable_scope.variable_scope("root", partitioner=partitioner):
+      v0 = variable_scope.get_variable(
+          var_name, dtype=dtypes.float32, shape=var_shape)
+      v0_list = v0._get_variable_list()
+      v0_part = v0._get_partitions()
+      self.assertEqual(len(v0_list), expected_axis_shards)
+      self.assertAllEqual(v0_part, expected_partitions)
+
+  def testMinMaxVariablePartitioner(self):
+    with self.cached_session():
+      # Partitioning a variable of shape=[2048] with a minimum of 2K per slice.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=2 << 10,
+          var_name="v0_0",
+          var_shape=[2048],
+          expected_axis_shards=4,
+          expected_partitions=[4])
+
+      # Partitioning a variable of shape=[2048, 1024] with a minimum of 256K per
+      # slice.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=256 << 10,
+          var_name="v0",
+          var_shape=[2048, 1024],
+          expected_axis_shards=32,
+          expected_partitions=[32, 1])
+
+      # max_partitions restricts partitioning of the variable.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=16,
+          axis=0,
+          min_slice_size=256 << 10,
+          var_name="v1_max",
+          var_shape=[2048, 1024],
+          expected_axis_shards=16,
+          expected_partitions=[16, 1])
+      self._testMinMaxVariablePartitioner(
+          max_partitions=1,
+          axis=0,
+          min_slice_size=256 << 10,
+          var_name="v2_max",
+          var_shape=[2048, 1024],
+          expected_axis_shards=1,
+          expected_partitions=[1, 1])
+
+      # Reducing/Increasing min_slice_size proportionately increases/reduces the
+      # number of partitions.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=128 << 10,
+          var_name="v3_slice",
+          var_shape=[2048, 1024],
+          expected_axis_shards=64,
+          expected_partitions=[64, 1])
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=512 << 10,
+          var_name="v4_slice",
+          var_shape=[2048, 1024],
+          expected_axis_shards=16,
+          expected_partitions=[16, 1])
+
+      # Partitioning the variable along a different axis.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=1,
+          min_slice_size=256 << 10,
+          var_name="v5_axis",
+          var_shape=[64, 1024, 1, 3],
+          expected_axis_shards=3,
+          expected_partitions=[1, 3, 1, 1])
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=3,
+          min_slice_size=256 << 10,
+          var_name="v6_axis",
+          var_shape=[64, 1024, 1, 3],
+          expected_axis_shards=3,
+          expected_partitions=[1, 1, 1, 3])
+
+      # Can not partition the variable more than what its shape allows.
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=256 << 10,
+          var_name="v7_shape",
+          var_shape=[16, 128, 1024],
+          expected_axis_shards=16,
+          expected_partitions=[16, 1, 1])
+      self._testMinMaxVariablePartitioner(
+          max_partitions=100,
+          axis=0,
+          min_slice_size=256 << 10,
+          var_name="v8_shape",
+          var_shape=[4, 512, 1024],
+          expected_axis_shards=4,
+          expected_partitions=[4, 1, 1])
+
+
+def _IotaInitializer(shape, dtype=dtypes.float32, partition_info=None):
+  assert dtype == dtypes.float32
+  if len(shape) == 1:
+    return range(shape[0])
+  else:
+    val = _IotaInitializer(shape[1:], dtype)
+    return [[(10**i) * v for v in val] for i in range(shape[0])]
+
+
+class PartitionedVariablesTestCase(test.TestCase):
+
+  def _TestSaveSpec(self, slices, expected_specs):
+    self.assertEqual(len(expected_specs), len(slices))
+    for i in range(len(expected_specs)):
+      self.assertEqual(expected_specs[i], slices[i]._save_slice_info.spec)
+
+  def testVecConstantInit(self):
+    with self.cached_session():
+      rnd_par = constant_op.constant([1, 2, 3, 4])
+      vs = partitioned_variables.create_partitioned_variables([4], [4], rnd_par)
+      self.evaluate(variables.global_variables_initializer())
+      val = array_ops.concat(vs, 0)
+      rnd = self.evaluate(rnd_par)
+      self.assertAllClose(rnd, val)
+      self.assertEqual([dtypes.int32] * 4, [v.dtype.base_dtype for v in vs])
+      self._TestSaveSpec(vs, ["4 0,1", "4 1,1", "4 2,1", "4 3,1"])
+
+  def testConstantInit(self):
+    with self.cached_session():
+      rnd_par = constant_op.constant([[1, 2, 3, 4], [5, 6, 7, 8]])
+      vs = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
+                                                              rnd_par)
+      self.evaluate(variables.global_variables_initializer())
+      val = array_ops.concat(vs, 1)
+      rnd = self.evaluate(rnd_par)
+      self.assertAllClose(rnd, val)
+      self.assertEqual([dtypes.int32] * 2, [v.dtype.base_dtype for v in vs])
+      self._TestSaveSpec(vs, ["2 4 0,2:0,2", "2 4 0,2:2,2"])
+
+  def _testNameHelper(self, use_resource=False):
+    with self.cached_session():
+      rnd_par = constant_op.constant([[1, 2, 3, 4], [5, 6, 7, 8]])
+      with variable_scope.variable_scope("hi", use_resource=use_resource):
+        vs1 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
+                                                                 rnd_par)
+        vs2 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
+                                                                 rnd_par)
+      self.evaluate(variables.global_variables_initializer())
+      var1_name = vs1[0]._save_slice_info.full_name
+      var2_name = vs2[0]._save_slice_info.full_name
+      self.assertEqual("hi/PartitionedVariable", var1_name)
+      self.assertEqual("hi/PartitionedVariable_1", var2_name)
+      self.assertEqual(var1_name + "/part_0:0", vs1[0].name)
+      self.assertEqual(var1_name + "/part_1:0", vs1[1].name)
+      self.assertEqual(var2_name + "/part_0:0", vs2[0].name)
+      self.assertEqual(var2_name + "/part_1:0", vs2[1].name)
+    # Test same variable.
+    with self.cached_session():
+      rnd_par = constant_op.constant([[1, 2, 3, 4], [5, 6, 7, 8]])
+      with variable_scope.variable_scope(
+          "hola", use_resource=use_resource) as vs:
+        vs1 = partitioned_variables.create_partitioned_variables(
+            [2, 4], [1, 2], rnd_par, dtype=dtypes.int32)
+      with variable_scope.variable_scope(
+          vs, reuse=True, use_resource=use_resource):
+        vs2 = partitioned_variables.create_partitioned_variables(
+            [2, 4], [1, 2], rnd_par, dtype=dtypes.int32)
+      self.evaluate(variables.global_variables_initializer())
+      var1_name = vs1[0]._save_slice_info.full_name
+      var2_name = vs2[0]._save_slice_info.full_name
+      self.assertEqual("hola/PartitionedVariable", var1_name)
+      self.assertEqual("hola/PartitionedVariable", var2_name)
+      self.assertEqual(var1_name + "/part_0:0", vs1[0].name)
+      self.assertEqual(var1_name + "/part_1:0", vs1[1].name)
+      self.assertEqual(var2_name + "/part_0:0", vs2[0].name)
+      self.assertEqual(var2_name + "/part_1:0", vs2[1].name)
+    # Test name_scope
+    with self.cached_session():
+      rnd_par = constant_op.constant([[1, 2, 3, 4], [5, 6, 7, 8]])
+      with ops.name_scope("ola"):
+        vs1 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
+                                                                 rnd_par)
+        vs2 = partitioned_variables.create_partitioned_variables([2, 4], [1, 2],
+                                                                 rnd_par)
+      self.evaluate(variables.global_variables_initializer())
+      var1_name = vs1[0]._save_slice_info.full_name
+      var2_name = vs2[0]._save_slice_info.full_name
+      # Currently, the name scope 'ola' has no effect.
+      self.assertEqual("PartitionedVariable", var1_name)
+      self.assertEqual("PartitionedVariable_1", var2_name)
+      self.assertEqual(var1_name + "/part_0:0", vs1[0].name)
+      self.assertEqual(var1_name + "/part_1:0", vs1[1].name)
+      self.assertEqual(var2_name + "/part_0:0", vs2[0].name)
+      self.assertEqual(var2_name + "/part_1:0", vs2[1].name)
+
+  @test_util.run_deprecated_v1
+  def testName(self):
+    self._testNameHelper(use_resource=False)
+
+  def testResourceName(self):
+    self._testNameHelper(use_resource=True)
+
+  def testRandomInitValue(self):
+    with self.cached_session():
+      rnd = variables.Variable(random_ops.random_uniform([200, 40]))
+      vs = partitioned_variables.create_partitioned_variables(
+          rnd.get_shape(), [1, 10], initialized_value(rnd))
+      self.evaluate(variables.global_variables_initializer())
+      val = array_ops.concat(vs, 1)
+      rnd = self.evaluate(rnd)
+      self.assertAllClose(rnd, val)
+      self.assertEqual([dtypes.float32] * 10, [v.dtype.base_dtype for v in vs])
+      self._TestSaveSpec(vs, [
+          "200 40 0,200:0,4", "200 40 0,200:4,4", "200 40 0,200:8,4",
+          "200 40 0,200:12,4", "200 40 0,200:16,4", "200 40 0,200:20,4",
+          "200 40 0,200:24,4", "200 40 0,200:28,4", "200 40 0,200:32,4",
+          "200 40 0,200:36,4"
+      ])
+
+  def testRandomInitUnevenPartitions(self):
+    with self.cached_session():
+      rnd = variables.Variable(
+          random_ops.random_uniform([20, 43], dtype=dtypes.float64))
+      var_lists = [
+          partitioned_variables.create_partitioned_variables(
+              rnd.get_shape(), [1, i], initialized_value(rnd))
+          for i in range(1, 10)
+      ]
+      self.evaluate(variables.global_variables_initializer())
+      rnd_val = self.evaluate(rnd)
+      # Only check the slice save specs for the first 5 tf.
+      save_specs = [
+          # One slice
+          ["20 43 0,20:0,43"],
+          # Two slices
+          ["20 43 0,20:0,22", "20 43 0,20:22,21"],
+          # Three slices
+          ["20 43 0,20:0,15", "20 43 0,20:15,14", "20 43 0,20:29,14"],
+          # Four slices
+          [
+              "20 43 0,20:0,11", "20 43 0,20:11,11", "20 43 0,20:22,11",
+              "20 43 0,20:33,10"
+          ],
+          # Five slices
+          [
+              "20 43 0,20:0,9", "20 43 0,20:9,9", "20 43 0,20:18,9",
+              "20 43 0,20:27,8", "20 43 0,20:35,8"
+          ]
+      ]
+      for i, vs in enumerate(var_lists):
+        var_val = array_ops.concat(vs, 1)
+        self.assertAllClose(rnd_val, var_val)
+        self.assertEqual([dtypes.float64] * len(vs),
+                         [v.dtype.base_dtype for v in vs])
+        if i < len(save_specs):
+          self._TestSaveSpec(vs, save_specs[i])
+
+  def testDegenerate(self):
+    with self.cached_session():
+      rnd = variables.Variable(random_ops.random_uniform([10, 43]))
+      vs = partitioned_variables.create_partitioned_variables(
+          rnd.get_shape(), [1, 1], initialized_value(rnd))
+      self.evaluate(variables.global_variables_initializer())
+      val = array_ops.concat(vs, 0)
+      rnd = self.evaluate(rnd)
+      self.assertAllClose(rnd, val)
+      self._TestSaveSpec(vs, ["10 43 0,10:0,43"])
+
+  def testSliceSizeOne(self):
+    with self.cached_session():
+      rnd = variables.Variable(random_ops.random_uniform([10, 43]))
+      vs = partitioned_variables.create_partitioned_variables(
+          rnd.get_shape(), [10, 1], initialized_value(rnd))
+      self.evaluate(variables.global_variables_initializer())
+      val = array_ops.concat(vs, 0)
+      rnd = self.evaluate(rnd)
+      self.assertAllClose(rnd, val)
+      self._TestSaveSpec(vs, [
+          "10 43 0,1:0,43", "10 43 1,1:0,43", "10 43 2,1:0,43",
+          "10 43 3,1:0,43", "10 43 4,1:0,43", "10 43 5,1:0,43",
+          "10 43 6,1:0,43", "10 43 7,1:0,43", "10 43 8,1:0,43", "10 43 9,1:0,43"
+      ])
+
+  def testIotaInitializer(self):
+    self.assertAllClose([0., 1., 2., 3.], _IotaInitializer([4]))
+    self.assertAllClose([[0., 1.], [0., 10.], [0., 100.], [0., 1000.]],
+                        _IotaInitializer([4, 2]))
+    with self.cached_session():
+      vs = partitioned_variables.create_partitioned_variables([13, 5], [3, 1],
+                                                              _IotaInitializer)
+      self.evaluate(variables.global_variables_initializer())
+      slice0 = _IotaInitializer([5, 5])
+      slice1 = _IotaInitializer([4, 5])
+      slice2 = _IotaInitializer([4, 5])
+      val = array_ops.concat(vs, 0)
+      self.assertAllClose(slice0 + slice1 + slice2, val)
+      self._TestSaveSpec(vs, ["13 5 0,5:0,5", "13 5 5,4:0,5", "13 5 9,4:0,5"])
+
+  @test_util.run_deprecated_v1
+  def testRandomInitializer(self):
+    # Sanity check that the slices uses a different seed when using a random
+    # initializer function.
+    with self.cached_session():
+      var0, var1 = partitioned_variables.create_partitioned_variables(
+          [20, 12], [1, 2], init_ops.random_uniform_initializer())
+      self.evaluate(variables.global_variables_initializer())
+      val0, val1 = self.evaluate(var0).flatten(), self.evaluate(var1).flatten()
+      self.assertTrue(np.linalg.norm(val0 - val1) > 1e-6)
+    # Negative test that proves that slices have the same values if
+    # the random initializer uses a seed.
+    with self.cached_session():
+      var0, var1 = partitioned_variables.create_partitioned_variables(
+          [20, 12], [1, 2], init_ops.random_uniform_initializer(seed=201))
+      self.evaluate(variables.global_variables_initializer())
+      val0, val1 = self.evaluate(var0).flatten(), self.evaluate(var1).flatten()
+      self.assertAllClose(val0, val1)
+
+  def testSomeErrors(self):
+    with self.cached_session():
+      rnd = variables.Variable(random_ops.random_uniform([10, 43]))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10], [1, 1], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 20], [1], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 43], [1], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 43], [1, 2, 3], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 43], [11, 1], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 43], [20, 1], initialized_value(rnd))
+      with self.assertRaises(ValueError):
+        partitioned_variables.create_partitioned_variables(
+            [10, 43], [1, 50], initialized_value(rnd))
+
+  @test_util.run_deprecated_v1
+  def testControlDepsNone(self):
+    with self.cached_session() as session:
+      c = constant_op.constant(1.0)
+      with ops.control_dependencies([c]):
+        # d get the control dependency.
+        d = constant_op.constant(2.0)
+        # Partitioned variables do not.
+        var_x = variable_scope.get_variable(
+            "x",
+            shape=[2],
+            initializer=init_ops.ones_initializer(),
+            partitioner=partitioned_variables.variable_axis_size_partitioner(4))
+
+        ops_before_read = session.graph.get_operations()
+        var_x.as_tensor()  # Caches the ops for subsequent reads.
+        reading_ops = [
+            op for op in session.graph.get_operations()
+            if op not in ops_before_read
+        ]
+
+      self.assertEqual([c.op], d.op.control_inputs)
+      # Tests that no control dependencies are added to reading a partitioned
+      # variable which is similar to reading a variable.
+      for op in reading_ops:
+        self.assertEqual([], op.control_inputs)
+
+  @test_util.run_deprecated_v1
+  def testConcat(self):
+    with self.cached_session() as session:
+      var_x = variable_scope.get_variable(
+          "x",
+          initializer=constant_op.constant([1., 2.]),
+          partitioner=partitioned_variables.variable_axis_size_partitioner(4))
+
+      c = constant_op.constant(1.0)
+      with ops.control_dependencies([c]):
+        ops_before_concat = session.graph.get_operations()
+        value = var_x._concat()  # pylint: disable=protected-access
+        concat_ops = [
+            op for op in session.graph.get_operations()
+            if op not in ops_before_concat
+        ]
+
+      concat_control_inputs = [
+          ci for op in concat_ops for ci in op.control_inputs
+      ]
+      self.assertTrue(
+          c.op in concat_control_inputs,
+          "var_x._concat() should get control dependencies from its scope.")
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(value, var_x.as_tensor())
+
+  def testMetaGraphSaveLoad(self):
+    save_prefix = os.path.join(self.get_temp_dir(), "ckpt")
+    save_graph = ops.Graph()
+    with save_graph.as_default(), self.session(
+        graph=save_graph) as session:
+      partitioner = partitioned_variables.fixed_size_partitioner(5, axis=0)
+      with variable_scope.variable_scope("root", partitioner=partitioner):
+        v0 = variable_scope.get_variable(
+            "v0", dtype=dtypes.float32, shape=(10, 10))
+        v0_list = v0._get_variable_list()
+        v0_part = v0._get_partitions()
+        self.assertEqual(len(v0_list), 5)
+        self.assertAllEqual(v0_part, (5, 1))
+        self.evaluate(variables.global_variables_initializer())
+
+        save_graph.get_collection_ref("partvar").append(v0)
+        saver = saver_lib.Saver()
+        save_graph.finalize()
+        save_path = saver.save(sess=session, save_path=save_prefix)
+        previous_value = session.run(
+            save_graph.get_tensor_by_name(v0.name + ":0"))
+
+    restore_graph = ops.Graph()
+    with restore_graph.as_default(), self.session(
+        graph=restore_graph) as session:
+      saver = saver_lib.import_meta_graph(save_path + ".meta")
+      saver.restore(sess=session, save_path=save_path)
+      v0, = save_graph.get_collection_ref("partvar")
+      self.assertIsInstance(v0, variables.PartitionedVariable)
+      self.assertAllEqual(
+          previous_value,
+          session.run(restore_graph.get_tensor_by_name(v0.name + ":0")))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/variables/variable_ops_test.py

@@ -0,0 +1,299 @@
+# Copyright 2015 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 tensorflow.ops.tf.variable_op."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import gen_state_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+_NP_TO_TF = {
+    dtypes.bfloat16.as_numpy_dtype: dtypes.bfloat16,
+    np.float16: dtypes.float16,
+    np.float32: dtypes.float32,
+    np.float64: dtypes.float64,
+    np.int32: dtypes.int32,
+    np.int64: dtypes.int64,
+    np.complex64: dtypes.complex64,
+    np.complex128: dtypes.complex128,
+}
+
+
+class VariableOpTest(test.TestCase):
+
+  def _initFetch(self, x, tftype, use_gpu=None):
+    with self.test_session(use_gpu=use_gpu):
+      p = state_ops.variable_op(x.shape, tftype)
+      op = state_ops.assign(p, x)
+      op.op.run()
+      return self.evaluate(p)
+
+  def _testTypes(self, vals):
+    for dtype in [
+        np.float16, np.float32, np.float64, np.complex64, np.complex128,
+        np.int32, np.int64, dtypes.bfloat16.as_numpy_dtype
+    ]:
+      self.setUp()
+      x = vals.astype(dtype)
+      tftype = _NP_TO_TF[dtype]
+      self.assertAllEqual(x, self._initFetch(x, tftype, use_gpu=False))
+      # NOTE(touts): the GPU test should pass for all types, whether the
+      # Variable op has an implementation for that type on GPU as we expect
+      # that Variable and Assign have GPU implementations for matching tf.
+      self.assertAllEqual(x, self._initFetch(x, tftype, use_gpu=True))
+
+  @test_util.run_deprecated_v1
+  def testBasic(self):
+    self._testTypes(np.arange(0, 20).reshape([4, 5]))
+
+  @test_util.run_deprecated_v1
+  def testset_shape(self):
+    p = state_ops.variable_op([1, 2], dtypes.float32)
+    self.assertEqual([1, 2], p.get_shape())
+    p = state_ops.variable_op([1, 2], dtypes.float32, set_shape=False)
+    self.assertEqual(tensor_shape.unknown_shape(), p.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssign(self):
+    for dtype in [dtypes.float32, dtypes.int64, dtypes.uint32, dtypes.uint8]:
+      value = np.array([[42, 43]])
+      var = state_ops.variable_op(value.shape, dtype)
+      self.assertShapeEqual(value, var)
+      assigned = state_ops.assign(var, value)
+      self.assertShapeEqual(value, assigned)
+
+  @test_util.run_deprecated_v1
+  def testAssignNoValidateShape(self):
+    value = np.array([[42.0, 43.0]])
+    var = state_ops.variable_op(value.shape, dtypes.float32)
+    self.assertShapeEqual(value, var)
+    assigned = state_ops.assign(var, value, validate_shape=False)
+    self.assertShapeEqual(value, assigned)
+
+  @test_util.run_deprecated_v1
+  def testAssignNoVarShape(self):
+    value = np.array([[42.0, 43.0]])
+    var = state_ops.variable_op(value.shape, dtypes.float32, set_shape=False)
+    self.assertEqual(tensor_shape.unknown_shape(), var.get_shape())
+    assigned = state_ops.assign(var, value)
+    self.assertShapeEqual(value, assigned)
+
+  @test_util.run_deprecated_v1
+  def testAssignNoVarShapeNoValidateShape(self):
+    value = np.array([[42.0, 43.0]])
+    var = state_ops.variable_op(value.shape, dtypes.float32, set_shape=False)
+    self.assertEqual(tensor_shape.unknown_shape(), var.get_shape())
+    assigned = state_ops.assign(var, value, validate_shape=False)
+    self.assertShapeEqual(value, assigned)
+
+  def _NewShapelessTensor(self):
+    tensor = array_ops.placeholder(dtypes.float32)
+    self.assertEqual(tensor_shape.unknown_shape(), tensor.get_shape())
+    return tensor
+
+  @test_util.run_deprecated_v1
+  def testAssignNoValueShape(self):
+    value = self._NewShapelessTensor()
+    shape = [1, 2]
+    var = state_ops.variable_op(shape, dtypes.float32)
+    assigned = state_ops.assign(var, value)
+    self.assertEqual(shape, var.get_shape())
+    self.assertEqual(shape, assigned.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignNoValueShapeNoValidateShape(self):
+    value = self._NewShapelessTensor()
+    shape = [1, 2]
+    var = state_ops.variable_op(shape, dtypes.float32)
+    self.assertEqual(shape, var.get_shape())
+    assigned = state_ops.assign(var, value, validate_shape=False)
+    self.assertEqual(tensor_shape.unknown_shape(), assigned.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignNoShape(self):
+    with self.cached_session():
+      value = self._NewShapelessTensor()
+      var = state_ops.variable_op([1, 2], dtypes.float32, set_shape=False)
+      self.assertEqual(tensor_shape.unknown_shape(), var.get_shape())
+      self.assertEqual(tensor_shape.unknown_shape(),
+                       state_ops.assign(var, value).get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignNoShapeNoValidateShape(self):
+    with self.cached_session():
+      value = self._NewShapelessTensor()
+      var = state_ops.variable_op([1, 2], dtypes.float32, set_shape=False)
+      self.assertEqual(tensor_shape.unknown_shape(), var.get_shape())
+      self.assertEqual(
+          tensor_shape.unknown_shape(),
+          state_ops.assign(var, value, validate_shape=False).get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignUpdate(self):
+    for dtype in [
+        dtypes.float32, dtypes.int64, dtypes.uint32, dtypes.uint8,
+        dtypes.bfloat16
+    ]:
+      var = state_ops.variable_op([1, 2], dtype)
+      added = state_ops.assign_add(var, [[2, 3]])
+      self.assertEqual([1, 2], added.get_shape())
+      subbed = state_ops.assign_sub(var, [[12, 13]])
+      self.assertEqual([1, 2], subbed.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignUpdateNoVarShape(self):
+    var = state_ops.variable_op([1, 2], dtypes.float32, set_shape=False)
+    added = state_ops.assign_add(var, [[2.0, 3.0]])
+    self.assertEqual([1, 2], added.get_shape())
+    subbed = state_ops.assign_sub(var, [[12.0, 13.0]])
+    self.assertEqual([1, 2], subbed.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignUpdateNoValueShape(self):
+    var = state_ops.variable_op([1, 2], dtypes.float32)
+    added = state_ops.assign_add(var, self._NewShapelessTensor())
+    self.assertEqual([1, 2], added.get_shape())
+    subbed = state_ops.assign_sub(var, self._NewShapelessTensor())
+    self.assertEqual([1, 2], subbed.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testAssignUpdateNoShape(self):
+    var = state_ops.variable_op([1, 2], dtypes.float32, set_shape=False)
+    added = state_ops.assign_add(var, self._NewShapelessTensor())
+    self.assertEqual(tensor_shape.unknown_shape(), added.get_shape())
+    subbed = state_ops.assign_sub(var, self._NewShapelessTensor())
+    self.assertEqual(tensor_shape.unknown_shape(), subbed.get_shape())
+
+  @test_util.run_deprecated_v1
+  def testTemporaryVariable(self):
+    with test_util.use_gpu():
+      var = gen_state_ops.temporary_variable([1, 2],
+                                             dtypes.float32,
+                                             var_name="foo")
+      var = state_ops.assign(var, [[4.0, 5.0]])
+      var = state_ops.assign_add(var, [[6.0, 7.0]])
+      final = gen_state_ops.destroy_temporary_variable(var, var_name="foo")
+      self.assertAllClose([[10.0, 12.0]], self.evaluate(final))
+
+  @test_util.run_deprecated_v1
+  def testDestroyNonexistentTemporaryVariable(self):
+    with test_util.use_gpu():
+      var = gen_state_ops.temporary_variable([1, 2], dtypes.float32)
+      final = gen_state_ops.destroy_temporary_variable(var, var_name="bad")
+      with self.assertRaises(errors.NotFoundError):
+        self.evaluate(final)
+
+  @test_util.run_deprecated_v1
+  def testDuplicateTemporaryVariable(self):
+    with test_util.use_gpu():
+      var1 = gen_state_ops.temporary_variable([1, 2],
+                                              dtypes.float32,
+                                              var_name="dup")
+      var1 = state_ops.assign(var1, [[1.0, 2.0]])
+      var2 = gen_state_ops.temporary_variable([1, 2],
+                                              dtypes.float32,
+                                              var_name="dup")
+      var2 = state_ops.assign(var2, [[3.0, 4.0]])
+      final = var1 + var2
+      with self.assertRaises(errors.AlreadyExistsError):
+        self.evaluate(final)
+
+  @test_util.run_deprecated_v1
+  def testDestroyTemporaryVariableTwice(self):
+    with test_util.use_gpu():
+      var = gen_state_ops.temporary_variable([1, 2], dtypes.float32)
+      val1 = gen_state_ops.destroy_temporary_variable(var, var_name="dup")
+      val2 = gen_state_ops.destroy_temporary_variable(var, var_name="dup")
+      final = val1 + val2
+      with self.assertRaises(errors.NotFoundError):
+        self.evaluate(final)
+
+  @test_util.run_deprecated_v1
+  def testTemporaryVariableNoLeak(self):
+    with test_util.use_gpu():
+      var = gen_state_ops.temporary_variable([1, 2],
+                                             dtypes.float32,
+                                             var_name="bar")
+      final = array_ops.identity(var)
+      self.evaluate(final)
+
+  @test_util.run_deprecated_v1
+  def testTwoTemporaryVariablesNoLeaks(self):
+    with test_util.use_gpu():
+      var1 = gen_state_ops.temporary_variable([1, 2],
+                                              dtypes.float32,
+                                              var_name="var1")
+      var2 = gen_state_ops.temporary_variable([1, 2],
+                                              dtypes.float32,
+                                              var_name="var2")
+      final = var1 + var2
+      self.evaluate(final)
+
+  @test_util.run_deprecated_v1
+  def testAssignDependencyAcrossDevices(self):
+    with test_util.use_gpu():
+      # The variable and an op to increment it are on the GPU.
+      var = state_ops.variable_op([1], dtypes.float32)
+      self.evaluate(state_ops.assign(var, [1.0]))
+      increment = state_ops.assign_add(var, [1.0])
+      with ops.control_dependencies([increment]):
+        with test_util.force_cpu():
+          # This mul op is pinned to the CPU, but reads the variable from the
+          # GPU. The test ensures that the dependency on 'increment' is still
+          # honored, i.e., the Send and Recv from GPU to CPU should take place
+          # only after the increment.
+          result = math_ops.multiply(var, var)
+      self.assertAllClose([4.0], self.evaluate(result))
+
+  @test_util.run_deprecated_v1
+  def testIsVariableInitialized(self):
+    for use_gpu in [True, False]:
+      with self.test_session(use_gpu=use_gpu):
+        v0 = state_ops.variable_op([1, 2], dtypes.float32)
+        self.assertEqual(False, variable_v1.is_variable_initialized(v0).eval())
+        state_ops.assign(v0, [[2.0, 3.0]]).eval()
+        self.assertEqual(True, variable_v1.is_variable_initialized(v0).eval())
+
+
+  @test_util.run_deprecated_v1
+  def testString(self):
+    data = array_ops_stack.stack([b"data"])
+    buffer_var = variable_v1.VariableV1(
+        initial_value=array_ops.zeros(shape=(), dtype=dtypes.string),
+        trainable=False,
+        collections=[ops.GraphKeys.LOCAL_VARIABLES],
+        name="buffer",
+        dtype=dtypes.string,
+        validate_shape=False,
+        use_resource=False)
+    result = state_ops.assign(buffer_var, data, validate_shape=False)
+    with self.cached_session() as sess:
+      sess.run(variables.local_variables_initializer())
+      self.assertEqual(result.eval(), b"data")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/kernel_tests/variables/variables_test.py

@@ -0,0 +1,962 @@
+# Copyright 2015 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 tf.py."""
+
+import functools
+import operator
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond as tf_cond
+from tensorflow.python.ops import gen_state_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.util import compat
+
+
+def initialized_value(var):
+  return tf_cond.cond(
+      variable_v1.is_variable_initialized(var), var.read_value,
+      lambda: var.initial_value)
+
+
+class VariablesTestCase(test.TestCase, parameterized.TestCase):
+
+  @test_util.run_deprecated_v1
+  def testDistributeStrategy(self):
+    v = variable_v1.VariableV1(0.0)
+    self.assertIsNone(v._distribute_strategy)
+
+  @test_util.run_v1_only("b/120545219")
+  def testInitialization(self):
+    with self.cached_session():
+      var0 = variable_v1.VariableV1(0.0)
+      self.assertEqual("Variable:0", var0.name)
+      self.assertEqual("Variable", var0._shared_name)
+      self.assertEqual([], var0.get_shape())
+      self.assertEqual([], var0.get_shape())
+      self.assertEqual([], var0.shape)
+
+      var1 = variable_v1.VariableV1(1.1)
+      self.assertEqual("Variable_1:0", var1.name)
+      self.assertEqual("Variable_1", var1._shared_name)
+      self.assertEqual([], var1.get_shape())
+      self.assertEqual([], var1.get_shape())
+      self.assertEqual([], var1.shape)
+
+      with self.assertRaisesOpError("Attempting to use uninitialized value"):
+        self.evaluate(var0)
+
+      with self.assertRaisesOpError("Attempting to use uninitialized value"):
+        self.evaluate(var1)
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose(0.0, self.evaluate(var0))
+      self.assertAllClose(1.1, self.evaluate(var1))
+
+  @test_util.run_v1_only("b/120545219")
+  def testInitializationOrder(self):
+    with self.cached_session():
+      rnd = variables.Variable(random_ops.random_uniform([3, 6]), name="rnd")
+      self.assertEqual("rnd:0", rnd.name)
+      self.assertEqual([3, 6], rnd.get_shape())
+      self.assertEqual([3, 6], rnd.get_shape())
+      self.assertEqual([3, 6], rnd.shape)
+
+      dep = variables.Variable(initialized_value(rnd), name="dep")
+      self.assertEqual("dep:0", dep.name)
+      self.assertEqual([3, 6], dep.get_shape())
+      self.assertEqual([3, 6], dep.get_shape())
+      self.assertEqual([3, 6], dep.shape)
+
+      # Currently have to set the shape manually for Add.
+      added_val = initialized_value(rnd) + initialized_value(dep) + 2.0
+      added_val.set_shape(rnd.get_shape())
+
+      depdep = variables.Variable(added_val, name="depdep")
+      self.assertEqual("depdep:0", depdep.name)
+      self.assertEqual([3, 6], depdep.get_shape())
+      self.assertEqual([3, 6], depdep.get_shape())
+      self.assertEqual([3, 6], depdep.shape)
+
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose(self.evaluate(rnd), self.evaluate(dep))
+      self.assertAllClose(
+          self.evaluate(rnd) + self.evaluate(dep) + 2.0, self.evaluate(depdep))
+
+  @test_util.run_deprecated_v1
+  def testCyclicInitializer(self):
+    with self.cached_session():
+      cyclic = while_loop.while_loop(
+          cond=lambda i: i < 10,
+          body=lambda i: i + 1,
+          loop_vars=(constant_op.constant(0),))
+      initial_value = variables._try_guard_against_uninitialized_dependencies(
+          "test", cyclic)
+      self.assertIs(initial_value, cyclic)
+
+  @test_util.run_deprecated_v1
+  def testIterableV1(self):
+    with self.assertRaisesRegex(TypeError, "not allowed.*Graph mode"):
+      for _ in variables.Variable(0.0):
+        pass
+    with self.assertRaisesRegex(TypeError, "not allowed.*Graph mode"):
+      for _ in variables.Variable([0.0, 1.0]):
+        pass
+
+  @test_util.run_v2_only
+  def testIterableV2(self):
+    with self.assertRaisesRegex(TypeError, "scalar tensor"):
+      for _ in variables.Variable(0.0):
+        pass
+    values = []
+    for v in variables.Variable([0.0, 1.0]):
+      values.append(v)
+    self.assertAllClose([0., 1.], values)
+
+  @test_util.run_deprecated_v1
+  def testAssignments(self):
+    with self.cached_session():
+      var = variables.Variable(0.0)
+      plus_one = var.assign_add(1.0)
+      minus_one = var.assign_sub(2.0)
+      four = var.assign(4.0)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(0.0, self.evaluate(var))
+
+      self.assertAllClose(1.0, self.evaluate(plus_one))
+      self.assertAllClose(1.0, self.evaluate(var))
+
+      self.assertAllClose(-1.0, self.evaluate(minus_one))
+      self.assertAllClose(-1.0, self.evaluate(var))
+
+      self.assertAllClose(4.0, self.evaluate(four))
+      self.assertAllClose(4.0, self.evaluate(var))
+
+  @test_util.run_deprecated_v1
+  def testResourceAssignments(self):
+    with self.session():
+      var = resource_variable_ops.ResourceVariable(0.0)
+      plus_one = var.assign_add(1.0)
+      minus_one = var.assign_sub(2.0)
+      four = var.assign(4.0)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(0.0, self.evaluate(var))
+
+      self.evaluate(plus_one)
+      self.assertAllClose(1.0, self.evaluate(var))
+
+      self.evaluate(minus_one)
+      self.assertAllClose(-1.0, self.evaluate(var))
+
+      self.evaluate(four)
+      self.assertAllClose(4.0, self.evaluate(var))
+
+  def testAssignDifferentShapesEagerNotAllowed(self):
+    with context.eager_mode():
+      var = variables.Variable(np.zeros(shape=[1, 1]))
+      with self.assertRaisesRegex(ValueError, "shape.*and.*are incompatible"):
+        var.assign(np.zeros(shape=[2, 2]))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAssignDifferentShapesAllowed(self):
+    var = variables.Variable(np.zeros(shape=[1, 1]),
+                             shape=tensor_shape.TensorShape(None))
+    self.evaluate(variables.global_variables_initializer())
+    self.assertAllEqual(np.zeros(shape=[1, 1]), var.read_value())
+    self.evaluate(var.assign(np.zeros(shape=[2, 2])))
+    self.assertAllEqual(np.zeros(shape=[2, 2]), var.read_value())
+
+  def testZeroSizeStringAssign(self):
+    with self.cached_session() as sess:
+      array = variable_v1.VariableV1(
+          initial_value=array_ops.zeros((0,), dtype=dtypes.string),
+          name="foo",
+          trainable=False,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES])
+      self.evaluate(variables.local_variables_initializer())
+      old_value = array.value()
+      copy_op = array.assign(old_value)
+      self.assertEqual([], list(self.evaluate(copy_op)))
+
+  def _countUpToTest(self, dtype):
+    with self.cached_session():
+      zero = constant_op.constant(0, dtype=dtype)
+      var = variables.Variable(zero)
+      count_up_to = var.count_up_to(3)
+
+      self.evaluate(variables.global_variables_initializer())
+      self.assertEqual(0, self.evaluate(var))
+
+      self.assertEqual(0, self.evaluate(count_up_to))
+      self.assertEqual(1, self.evaluate(var))
+
+      self.assertEqual(1, self.evaluate(count_up_to))
+      self.assertEqual(2, self.evaluate(var))
+
+      self.assertEqual(2, self.evaluate(count_up_to))
+      self.assertEqual(3, self.evaluate(var))
+
+      with self.assertRaisesOpError("Reached limit of 3"):
+        self.evaluate(count_up_to)
+      self.assertEqual(3, self.evaluate(var))
+
+      with self.assertRaisesOpError("Reached limit of 3"):
+        self.evaluate(count_up_to)
+      self.assertEqual(3, self.evaluate(var))
+
+  @test_util.run_deprecated_v1
+  def testCountUpToInt32(self):
+    self._countUpToTest(dtypes.int32)
+
+  @test_util.run_deprecated_v1
+  def testCountUpToInt64(self):
+    self._countUpToTest(dtypes.int64)
+
+  @test_util.run_v1_only("b/120545219")
+  def testControlDepsNone(self):
+    with self.cached_session():
+      c = constant_op.constant(1.0)
+      with ops.control_dependencies([c]):
+        # d get the control dep.
+        d = constant_op.constant(2.0)
+        # variables do not.
+        var_x = variable_v1.VariableV1(2.0)
+      self.assertEqual([c.op], d.op.control_inputs)
+      self.assertEqual([], var_x.initializer.control_inputs)
+      self.assertEqual([], var_x.value().op.control_inputs)
+      self.assertEqual([], var_x._ref().op.control_inputs)  # pylint: disable=protected-access
+
+  @test_util.run_v1_only("b/120545219")
+  def testControlFlow(self):
+    with self.cached_session() as sess:
+      v0 = variables.Variable(0, name="v0")
+      var_dict = {}
+
+      # Call get_variable in each of the cond clauses.
+      def var_in_then_clause():
+        v1 = variables.Variable(1, name="v1")
+        var_dict["v1"] = v1
+        return v1 + v0
+
+      def var_in_else_clause():
+        v2 = variables.Variable(2, name="v2")
+        var_dict["v2"] = v2
+        return v2 + v0
+
+      add = tf_cond.cond(
+          math_ops.less(v0, 10), var_in_then_clause, var_in_else_clause)
+      v1 = var_dict["v1"]
+      v2 = var_dict["v2"]
+      # We should be able to initialize and run v1 and v2 without initializing
+      # v0, even if the variable was created with a control dep on v0.
+      self.evaluate(v1.initializer)
+      self.assertEqual([1], self.evaluate(v1))
+      self.evaluate(v2.initializer)
+      self.assertEqual([2], self.evaluate(v2))
+      # v0 should still be uninitialized.
+      with self.assertRaisesRegex(errors_impl.OpError, "uninitialized"):
+        self.evaluate(v0)
+      # We should not be able to run 'add' yet.
+      with self.assertRaisesRegex(errors_impl.OpError, "uninitialized"):
+        self.evaluate(add)
+      # If we initialize v0 we should be able to run 'add'.
+      self.evaluate(v0.initializer)
+      self.evaluate(add)
+
+  @test_util.run_v1_only("b/120545219")
+  def testControlFlowInitialization(self):
+    """Expects an error if an initializer is in a control-flow scope."""
+    def cond(i, _):
+      return i < 10
+
+    def body(i, _):
+      zero = array_ops.zeros([], dtype=dtypes.int32)
+      v = variables.Variable(initial_value=zero)
+      return (i + 1, v.read_value())
+
+    with self.assertRaisesRegex(ValueError, "inside a control-flow"):
+      while_loop.while_loop(cond, body, [0, 0])
+
+  @test_util.run_deprecated_v1
+  def testUseVariableAsTensor(self):
+    with self.cached_session():
+      var_x = variables.Variable(2.0)
+      var_y = variables.Variable(3.0)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(2.0, self.evaluate(var_x))
+      self.assertAllClose(3.0, self.evaluate(var_y))
+      self.assertAllClose(5.0, self.evaluate(math_ops.add(var_x, var_y)))
+
+  @test_util.run_deprecated_v1
+  def testZeroSizeVarSameAsConst(self):
+    with self.cached_session():
+      zero_size_var = variables.Variable(array_ops.zeros([0, 2]))
+      zero_size_const = array_ops.ones([2, 0])
+      variable_mul = math_ops.matmul(zero_size_const, zero_size_var)
+      const_mul = math_ops.matmul(
+          zero_size_const, zero_size_const, transpose_b=True)
+      self.evaluate(variables.global_variables_initializer())
+      variable_output = self.evaluate(variable_mul)
+      self.assertAllClose(self.evaluate(const_mul), variable_output)
+      self.assertAllClose([[0., 0.], [0., 0.]], variable_output)
+
+  @test_util.run_deprecated_v1
+  def testCachingDevice(self):
+    with self.cached_session():
+      var = variables.Variable(2.0)
+      self.assertEqual(var.device, initialized_value(var).device)
+
+      var_cached = variables.Variable(2.0, caching_device="/job:foo")
+      self.assertFalse(var_cached.device.startswith("/job:foo"))
+      self.assertTrue(var_cached.value().device.startswith("/job:foo"))
+
+  @test_util.run_deprecated_v1
+  def testCollections(self):
+    with self.cached_session():
+      var_x = variable_v1.VariableV1(2.0)
+      var_y = variable_v1.VariableV1(2.0, trainable=False)
+      var_z = variable_v1.VariableV1(2.0, trainable=True)
+      var_t = variable_v1.VariableV1(
+          2.0,
+          trainable=True,
+          collections=[
+              ops.GraphKeys.TRAINABLE_VARIABLES, ops.GraphKeys.GLOBAL_VARIABLES
+          ])
+      self.assertEqual([var_x, var_y, var_z, var_t],
+                       variables.global_variables())
+      self.assertEqual([var_x, var_z, var_t], variables.trainable_variables())
+
+  @test_util.run_deprecated_v1
+  def testCollectionsWithScope(self):
+    with self.cached_session():
+      with ops.name_scope("scope_1"):
+        var_x = variable_v1.VariableV1(2.0)
+      with ops.name_scope("scope_2"):
+        var_y = variable_v1.VariableV1(2.0)
+
+      self.assertEqual([var_x, var_y], variables.global_variables())
+      self.assertEqual([var_x], variables.global_variables("scope_1"))
+      self.assertEqual([var_y], variables.global_variables("scope_2"))
+
+      self.assertEqual([var_x, var_y], variables.trainable_variables())
+      self.assertEqual([var_x], variables.trainable_variables("scope_1"))
+      self.assertEqual([var_y], variables.trainable_variables("scope_2"))
+
+  def testOperatorWrapping(self):
+    for attr in functools.WRAPPER_ASSIGNMENTS:
+      self.assertEqual(
+          getattr(variables.Variable.__add__, attr),
+          getattr(tensor.Tensor.__add__, attr))
+
+  @test_util.run_deprecated_v1
+  def testOperators(self):
+    with self.cached_session():
+      var_f = variables.Variable([2.0])
+      add = var_f + 0.0
+      radd = 1.0 + var_f
+      sub = var_f - 1.0
+      rsub = 1.0 - var_f
+      mul = var_f * 10.0
+      rmul = 10.0 * var_f
+      div = var_f / 10.0
+      rdiv = 10.0 / var_f
+      lt = var_f < 3.0
+      rlt = 3.0 < var_f
+      le = var_f <= 2.0
+      rle = 2.0 <= var_f
+      gt = var_f > 3.0
+      rgt = 3.0 > var_f
+      ge = var_f >= 2.0
+      rge = 2.0 >= var_f
+      neg = -var_f
+      abs_v = abs(var_f)
+
+      var_i = variables.Variable([20])
+      mod = var_i % 7
+      rmod = 103 % var_i
+
+      var_b = variables.Variable([True, False])
+      and_v = operator.and_(var_b, [True, True])
+      or_v = operator.or_(var_b, [False, True])
+      xor_v = operator.xor(var_b, [False, False])
+      invert_v = ~var_b
+
+      rnd = np.random.rand(4, 4).astype("f")
+      var_t = variables.Variable(rnd)
+      slice_v = var_t[2, 0:0]
+
+      var_m = variables.Variable([[2.0, 3.0]])
+      matmul = var_m.__matmul__([[10.0], [20.0]])
+      rmatmul = var_m.__rmatmul__([[10.0], [20.0]])
+
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([2.0], self.evaluate(add))
+      self.assertAllClose([3.0], self.evaluate(radd))
+      self.assertAllClose([1.0], self.evaluate(sub))
+      self.assertAllClose([-1.0], self.evaluate(rsub))
+      self.assertAllClose([20.0], self.evaluate(mul))
+      self.assertAllClose([20.0], self.evaluate(rmul))
+      self.assertAllClose([0.2], self.evaluate(div))
+      self.assertAllClose([5.0], self.evaluate(rdiv))
+      self.assertAllClose([-2.0], self.evaluate(neg))
+      self.assertAllClose([2.0], self.evaluate(abs_v))
+      self.assertAllClose([True], self.evaluate(lt))
+      self.assertAllClose([False], self.evaluate(rlt))
+      self.assertAllClose([True], self.evaluate(le))
+      self.assertAllClose([True], self.evaluate(rle))
+      self.assertAllClose([False], self.evaluate(gt))
+      self.assertAllClose([True], self.evaluate(rgt))
+      self.assertAllClose([True], self.evaluate(ge))
+      self.assertAllClose([True], self.evaluate(rge))
+
+      self.assertAllClose([6], self.evaluate(mod))
+      self.assertAllClose([3], self.evaluate(rmod))
+
+      self.assertAllClose([True, False], self.evaluate(and_v))
+      self.assertAllClose([True, True], self.evaluate(or_v))
+      self.assertAllClose([True, False], self.evaluate(xor_v))
+      self.assertAllClose([False, True], self.evaluate(invert_v))
+
+      self.assertAllClose(rnd[2, 0:0], self.evaluate(slice_v))
+
+      self.assertAllClose([[80.0]], self.evaluate(matmul))
+      self.assertAllClose([[20.0, 30.0], [40.0, 60.0]], self.evaluate(rmatmul))
+
+  @test_util.run_deprecated_v1
+  def testSession(self):
+    with self.cached_session() as sess:
+      var = variables.Variable([1, 12])
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose([1, 12], self.evaluate(var))
+
+  @test_util.run_v1_only("b/120545219")
+  def testColocation(self):
+    with ops.device("/job:ps"):
+      var = variable_v1.VariableV1(0, name="v")
+    with ops.device("/job:worker/task:7"):
+      assign_op = var.assign(1)
+    self.assertDeviceEqual("/job:ps", assign_op.device)
+    self.assertEqual([b"loc:@v"], assign_op.op.colocation_groups())
+
+  @test_util.run_v1_only("b/120545219")
+  def testInitializerFunction(self):
+    value = [[-42], [133.7]]
+    shape = [2, 1]
+    with self.cached_session():
+      initializer = lambda: constant_op.constant(value)
+
+      v1 = variables.Variable(initializer, dtype=dtypes.float32)
+      self.assertEqual(shape, v1.get_shape())
+      self.assertEqual(shape, v1.shape)
+      self.assertAllClose(value, self.evaluate(v1.initial_value))
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        self.evaluate(v1)
+
+      v2 = variables.Variable(
+          math_ops.negative(initialized_value(v1)), dtype=dtypes.float32)
+      self.assertEqual(v1.get_shape(), v2.get_shape())
+      self.assertEqual(v1.shape, v2.shape)
+      self.assertAllClose(np.negative(value), self.evaluate(v2.initial_value))
+
+      with self.assertRaises(errors_impl.FailedPreconditionError):
+        self.evaluate(v2)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllClose(np.negative(value), self.evaluate(v2))
+
+  def testConstraintArg(self):
+    constraint = lambda x: x
+    v = variables.Variable(
+        lambda: constant_op.constant(1.),
+        constraint=constraint)
+    self.assertEqual(v.constraint, constraint)
+
+    constraint = 0
+    with self.assertRaises(ValueError):
+      v = variables.Variable(
+          lambda: constant_op.constant(1.),
+          constraint=constraint)
+
+  @test_util.run_v1_only("b/120545219")
+  def testNoRefDataRace(self):
+    with self.cached_session():
+      a = variables.Variable([1, 2, 3], dtype=dtypes.float32)
+      b = variables.Variable(initialized_value(a) + 2)
+      c = variables.Variable(initialized_value(b) + 2)
+      self.evaluate(variables.global_variables_initializer())
+      self.assertAllEqual(self.evaluate(a), [1, 2, 3])
+      self.assertAllEqual(self.evaluate(b), [3, 4, 5])
+      self.assertAllEqual(self.evaluate(c), [5, 6, 7])
+
+  @test_util.run_deprecated_v1
+  def testInitializerFunctionDevicePlacement(self):
+    with self.cached_session():
+      initializer = lambda: constant_op.constant(42.0)
+      with ops.device("/cpu:100"):
+        v1 = variables.Variable(initializer, dtype=dtypes.float32, name="v1")
+      expected_device = "/device:CPU:100"
+      expected_group_v1 = [b"loc:@v1"]
+      self.assertEqual(expected_device, v1.op.device)
+      self.assertEqual(expected_group_v1, v1.op.colocation_groups())
+      for i in v1.initializer.inputs:
+        self.assertEqual(expected_group_v1, i.op.colocation_groups())
+
+      v2 = variables.Variable(initializer, dtype=dtypes.float32, name="v2")
+      expected_group_v2 = [b"loc:@v2"]
+      self.assertEqual(expected_group_v2, v2.op.colocation_groups())
+      for i in v2.initializer.inputs:
+        self.assertEqual(expected_group_v2, i.op.colocation_groups())
+
+  @test_util.run_v1_only("b/120545219")
+  def testVariableDefInitializedInstances(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v_def = variables.Variable(
+          initial_value=constant_op.constant(3.0)).to_proto()
+
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      # v describes a VariableDef-based variable without an initial value.
+      v = variables.Variable(variable_def=v_def)
+      self.assertEqual(3.0, self.evaluate(initialized_value(v)))
+
+      # initialized_value should not rerun the initializer_op if the variable
+      # has already been initialized elsewhere.
+      self.evaluate(v.assign(1.0))
+      self.assertEqual(1.0, self.evaluate(initialized_value(v)))
+
+    v_def.ClearField("initial_value_name")
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      # Restoring a legacy VariableDef proto that does not have
+      # initial_value_name set should still work.
+      v = variables.Variable(variable_def=v_def)
+      # We should also be able to re-export the variable to a new meta graph.
+      self.assertProtoEquals(v_def, v.to_proto())
+      # But attempts to use initialized_value will result in errors.
+      with self.assertRaises(ValueError):
+        self.evaluate(initialized_value(v))
+
+  def testTrainableInProto(self):
+    with ops.Graph().as_default():
+      non_trainable_variable = variables.Variable(
+          trainable=False,
+          initial_value=constant_op.constant(10.0))
+      self.assertEqual(
+          False,
+          variables.Variable(variable_def=non_trainable_variable.to_proto())
+          .trainable)
+      trainable_variable = variables.Variable(
+          trainable=True,
+          initial_value=constant_op.constant(10.0))
+      self.assertEqual(
+          True,
+          variables.Variable(variable_def=trainable_variable.to_proto())
+          .trainable)
+
+  def testSynchronizationAndAggregationSaved(self):
+    with ops.Graph().as_default():
+      original_variable = variables.Variable(
+          initial_value=constant_op.constant(10.0),
+          synchronization=variables.VariableSynchronization.NONE,
+          aggregation=variables.VariableAggregationV2.ONLY_FIRST_REPLICA)
+      self.assertEqual(variables.VariableSynchronization.NONE,
+                       original_variable.synchronization)
+      self.assertEqual(variables.VariableAggregation.ONLY_FIRST_REPLICA,
+                       original_variable.aggregation)
+
+      laundered = variables.Variable(
+          variable_def=original_variable.to_proto())
+      self.assertEqual(
+          variables.VariableSynchronization.NONE,
+          laundered.synchronization)
+      self.assertEqual(variables.VariableAggregationV2.ONLY_FIRST_REPLICA,
+                       laundered.aggregation)
+
+  @test_util.run_deprecated_v1
+  def testLoad(self):
+    with self.cached_session():
+      var = variables.Variable(np.zeros((5, 5), np.float32))
+      self.evaluate(variables.global_variables_initializer())
+      var.load(np.ones((5, 5), np.float32))
+
+      self.assertAllClose(np.ones((5, 5), np.float32), self.evaluate(var))
+
+  @test_util.run_v1_only("b/120545219")
+  def testRepr(self):
+    var = variable_v1.VariableV1(np.zeros((5, 5), np.float32), name="noop")
+    self.assertEqual(
+        "<tf.Variable 'noop:0' shape=(5, 5) dtype=float32_ref>",
+        repr(var))
+
+  def testVariableNamesPreserveNameScopesWithFunction(self):
+    v = None
+
+    @def_function.function
+    def create_variable():
+      with ops.name_scope("foo"):
+        nonlocal v
+        if v is None:
+          v = variables.Variable(0.0, name="bar")
+      self.assertEqual(v.name, "foo/bar:0")
+    with ops.get_default_graph().as_default():
+      create_variable()
+
+  @parameterized.parameters(variable_v1.VariableV1, variables.Variable)
+  def testTrainableVariable(self, cls):
+    v1 = cls(1.0)
+    self.assertEqual(True, v1.trainable)
+
+    v2 = cls(1.0, synchronization=variables.VariableSynchronization.ON_READ)
+    self.assertEqual(False, v2.trainable)
+
+    v3 = cls(1.0, synchronization=variables.VariableSynchronization.ON_READ,
+             trainable=True)
+    self.assertEqual(True, v3.trainable)
+
+    v4 = cls(1.0, synchronization=variables.VariableSynchronization.ON_READ,
+             trainable=False)
+    self.assertEqual(False, v4.trainable)
+
+  def testSaveSliceInfoFromSpecPasses(self):
+    save_slice_info = variables.Variable.SaveSliceInfo(
+        full_name="foo",
+        full_shape=[2, 3, 4],
+        var_offset=[0, 2, 0],
+        var_shape=[1, 1, 3])
+
+    save_slice_info_from_spec = variables.Variable.SaveSliceInfo.from_spec(
+        save_slice_info.spec)
+
+    self.assertEqual(save_slice_info.spec, save_slice_info_from_spec.spec)
+
+  @parameterized.parameters(
+      dict(spec="0", error_message="contain space-separated full_shape info"),
+      dict(spec="0:0", error_message="contain space-separated full_shape info"),
+      dict(spec="a b", error_message="full_shape must be a sequence of int"),
+      dict(spec="0 0:0", error_message="comma-separated pairs of offsets and"),
+      dict(spec="0 a,0:0:0", error_message="var_offset must be an integer"),
+      dict(spec="0 0,a:0:0", error_message="var_shape must be an integer"))
+  def testSaveSliceInfoFromSpecFails(self, spec, error_message):
+    with self.assertRaisesRegex(ValueError, error_message):
+      variables.Variable.SaveSliceInfo.from_spec(spec)
+
+
+class IsInitializedTest(test.TestCase):
+
+  def testNoVars(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      uninited = variables.report_uninitialized_variables()
+      self.assertEqual(0, self.evaluate(uninited).size)
+
+  def testAssertVariablesInitialized(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variables.Variable([1, 2], name="v")
+      w = variables.Variable([3, 4], name="w")
+      _ = v, w
+      uninited = variables.report_uninitialized_variables()
+      self.assertAllEqual(np.array([b"v", b"w"]), self.evaluate(uninited))
+      self.evaluate(variables.global_variables_initializer())
+      self.assertEqual(0, self.evaluate(uninited).size)
+
+  @test_util.run_v1_only("b/120545219")
+  def testVariableList(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variable_v1.VariableV1([1, 2], name="v")
+      w = variable_v1.VariableV1([3, 4], name="w")
+      uninited = variables.report_uninitialized_variables()
+      self.assertAllEqual(np.array([b"v", b"w"]), self.evaluate(uninited))
+      self.evaluate(w.initializer)
+      self.assertAllEqual(np.array([b"v"]), self.evaluate(uninited))
+      self.evaluate(v.initializer)
+      self.assertEqual(0, self.evaluate(uninited).size)
+
+  def testZeroSizeVarInitialized(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variables.Variable(array_ops.zeros([0, 2]), name="v")
+      uninited = variables.report_uninitialized_variables()
+      self.evaluate(v.initializer)  # not strictly necessary
+      self.assertEqual(0, self.evaluate(uninited).size)
+
+  def testTrainingWithZeroSizeVar(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      a = variables.Variable(array_ops.zeros([0, 2]))
+      b = variables.Variable(array_ops.ones([2, 2]))
+      objective = math_ops.reduce_sum(b + math_ops.matmul(
+          a, a, transpose_a=True))
+      self.evaluate(variables.global_variables_initializer())
+      do_opt = gradient_descent.GradientDescentOptimizer(0.1).minimize(
+          objective)
+      self.evaluate([do_opt])
+      self.assertAllClose([[0.9, 0.9], [0.9, 0.9]], self.evaluate(b))
+
+
+@test_util.run_v1_only("b/120545219")
+class ObsoleteIsInitializedTest(test.TestCase):
+
+  def testNoVars(self):
+    with ops.Graph().as_default():
+      self.assertEqual(None, variables.assert_variables_initialized())
+
+  def testVariables(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variable_v1.VariableV1([1, 2])
+      w = variable_v1.VariableV1([3, 4])
+      _ = v, w
+      inited = variables.assert_variables_initialized()
+      with self.assertRaisesOpError("Attempting to use uninitialized value"):
+        self.evaluate(inited)
+      self.evaluate(variables.global_variables_initializer())
+      self.evaluate(inited)
+
+  def testVariableList(self):
+    with ops.Graph().as_default(), self.cached_session() as sess:
+      v = variable_v1.VariableV1([1, 2])
+      w = variable_v1.VariableV1([3, 4])
+      inited = variables.assert_variables_initialized([v])
+      with self.assertRaisesOpError("Attempting to use uninitialized value"):
+        inited.op.run()
+      self.evaluate(w.initializer)
+      with self.assertRaisesOpError("Attempting to use uninitialized value"):
+        inited.op.run()
+      self.evaluate(v.initializer)
+      inited.op.run()
+
+
+class PartitionedVariableTest(test.TestCase):
+
+  def testPartitionedVariable(self):
+    with ops.Graph().as_default():
+      v0 = variables.Variable([0])
+      v1 = variables.Variable([1])
+      v0._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v0.name, [2], [0], [1]))
+      v1._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v0.name, [2], [1], [1]))
+      partitions = [2]
+
+      # Pass variable_list as [v1, v0] to ensure they are properly
+      # re-sorted to [v0, v1] based on their slice info offsets.
+      partitioned_variable = variables.PartitionedVariable(
+          name="two_vars",
+          shape=[2],
+          dtype=v0.dtype,
+          variable_list=[v1, v0],
+          partitions=partitions)
+
+      concatenated = ops.convert_to_tensor(partitioned_variable)
+      num_partitions = len(partitioned_variable)
+      iterated_partitions = list(partitioned_variable)
+      self.assertEqual(2, num_partitions)
+      self.assertEqual([v0, v1], iterated_partitions)
+      self.assertEqual([2], partitioned_variable.get_shape())
+      self.assertEqual([2], partitioned_variable.shape)
+      self.assertEqual([2], concatenated.get_shape())
+      self.assertEqual([2], concatenated.shape)
+
+  def testPartitionedVariableFailures(self):
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(ValueError, "empty"):
+        variables.PartitionedVariable(
+            name="fail",
+            shape=2,
+            dtype=dtypes.int32,
+            variable_list=[],
+            partitions=[])
+
+      with self.assertRaisesRegex(ValueError, "must have a save_slice_info"):
+        v0 = variables.Variable([0])
+        partitions = [1]
+        variables.PartitionedVariable(
+            name="two_vars",
+            shape=[1],
+            dtype=v0.dtype,
+            variable_list=[v0],
+            partitions=partitions)
+
+      with self.assertRaisesRegex(ValueError, "full shapes must match"):
+        v0 = variables.Variable([0])
+        v1 = variables.Variable([1])
+        v0._set_save_slice_info(
+            variables.Variable.SaveSliceInfo(v0.name, [2], [0], [1]))
+        v1._set_save_slice_info(
+            variables.Variable.SaveSliceInfo(v0.name, [2], [1], [1]))
+        partitions = [2]
+
+        variables.PartitionedVariable(
+            name="two_vars",
+            shape=[3],
+            dtype=v0.dtype,
+            variable_list=[v1, v0],
+            partitions=partitions)
+
+      with self.assertRaisesRegex(ValueError, "must be positive"):
+        v0 = variables.Variable([0])
+        v0._set_save_slice_info(
+            variables.Variable.SaveSliceInfo(v0.name, [2], [0], [1]))
+        partitions = [0]
+
+        variables.PartitionedVariable(
+            name="two_vars",
+            shape=[2],
+            dtype=v0.dtype,
+            variable_list=[v0],
+            partitions=partitions)
+
+  def testPartitionedVariableAssignments(self):
+    with ops.Graph().as_default(), self.cached_session():
+      v0 = variables.Variable(initial_value=[0.0])
+      v1 = variables.Variable(initial_value=[1.0])
+      v2 = variables.Variable(initial_value=[20.0])
+      v3 = variables.Variable(initial_value=[30.0])
+      v0._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v0.name, [2], [0], [1]))
+      v1._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v1.name, [2], [1], [1]))
+      v2._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v2.name, [2], [0], [1]))
+      v3._set_save_slice_info(
+          variables.Variable.SaveSliceInfo(v3.name, [2], [1], [1]))
+
+      partitions = [2]
+
+      # Pass variable_list as [v1, v0] to ensure they are properly
+      # re-sorted to [v0, v1] based on their slice info offsets.
+      pv_0 = variables.PartitionedVariable(
+          name="two_vars",
+          shape=[2],
+          dtype=v0.dtype,
+          variable_list=[v0, v1],
+          partitions=partitions)
+
+      pv_1 = variables.PartitionedVariable(
+          name="two_vars",
+          shape=[2],
+          dtype=v0.dtype,
+          variable_list=[v2, v3],
+          partitions=partitions)
+
+      deltas_a = constant_op.constant([1.0, 2.0])
+      deltas_b = constant_op.constant([3.0, 4.0])
+      ones = array_ops.ones([2])
+      plus_delta = pv_0.assign_add(deltas_a)
+      minus_delta = pv_0.assign_sub(deltas_b)
+      assign_ones = pv_0.assign(ones)
+
+      c_0 = constant_op.constant([2.0])
+      c_1 = constant_op.constant([3.0])
+      assign_list = pv_1.assign([c_0, c_1])
+      assign_part_value = pv_1.assign_add(assign_ones)
+      assign_part_var = pv_1.assign_sub(pv_0)
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertEqual([1.0], self.evaluate(plus_delta[0]))
+      self.assertEqual([1.0], self.evaluate(v0))
+      self.assertEqual([3.0], self.evaluate(plus_delta[1]))
+      self.assertEqual([3.0], self.evaluate(v1))
+
+      self.assertEqual([-2.0], self.evaluate(minus_delta[0]))
+      self.assertEqual([-2.0], self.evaluate(v0))
+      self.assertEqual([-1.0], self.evaluate(minus_delta[1]))
+      self.assertEqual([-1.0], self.evaluate(v1))
+
+      self.assertEqual([1.0], self.evaluate(assign_ones[0]))
+      self.assertEqual([1.0], self.evaluate(v0))
+      self.assertEqual([1.0], self.evaluate(assign_ones[1]))
+      self.assertEqual([1.0], self.evaluate(v1))
+
+      self.assertEqual([2.0], self.evaluate(assign_list[0]))
+      self.assertEqual([2.0], self.evaluate(v2))
+      self.assertEqual([3.0], self.evaluate(assign_list[1]))
+      self.assertEqual([3.0], self.evaluate(v3))
+
+      self.assertEqual([3.0], self.evaluate(assign_part_value[0]))
+      self.assertEqual([3.0], self.evaluate(v2))
+      self.assertEqual([4.0], self.evaluate(assign_part_value[1]))
+      self.assertEqual([4.0], self.evaluate(v3))
+
+      self.assertEqual([2.0], self.evaluate(assign_part_var[0]))
+      self.assertEqual([2.0], self.evaluate(v2))
+      self.assertEqual([3.0], self.evaluate(assign_part_var[1]))
+      self.assertEqual([3.0], self.evaluate(v3))
+
+
+class VariableContainerTest(test.TestCase):
+
+  def testContainer(self):
+    with ops.Graph().as_default():
+      v0 = variables.Variable([0])
+      with ops.container("l1"):
+        v1 = variables.Variable([1])
+        with ops.container("l2"):
+          v2 = variables.Variable([2])
+          special_v = gen_state_ops.variable(
+              shape=[1],
+              dtype=dtypes.float32,
+              name="VariableInL3",
+              container="l3",
+              shared_name="")
+        v3 = variables.Variable([3])
+      v4 = variables.Variable([4])
+    self.assertEqual(compat.as_bytes(""), v0.op.get_attr("container"))
+    self.assertEqual(compat.as_bytes("l1"), v1.op.get_attr("container"))
+    self.assertEqual(compat.as_bytes("l2"), v2.op.get_attr("container"))
+    self.assertEqual(compat.as_bytes("l3"), special_v.op.get_attr("container"))
+    self.assertEqual(compat.as_bytes("l1"), v3.op.get_attr("container"))
+    self.assertEqual(compat.as_bytes(""), v4.op.get_attr("container"))
+
+
+class AggregationModesTest(test.TestCase):
+
+  def testV1V2Equal(self):
+    v1 = variables.VariableAggregation
+    v2 = variables.VariableAggregationV2
+
+    self.assertEqual(v1.NONE, v2.NONE)
+    self.assertEqual(v1.SUM, v2.SUM)
+    self.assertEqual(v1.MEAN, v2.MEAN)
+    self.assertEqual(v1.ONLY_FIRST_REPLICA, v2.ONLY_FIRST_REPLICA)
+    self.assertEqual(v1.ONLY_FIRST_TOWER, v2.ONLY_FIRST_REPLICA)
+
+    self.assertEqual(v2.NONE, v1.NONE)
+    self.assertEqual(v2.SUM, v1.SUM)
+    self.assertEqual(v2.MEAN, v1.MEAN)
+    self.assertEqual(v2.ONLY_FIRST_REPLICA, v1.ONLY_FIRST_REPLICA)
+    self.assertEqual(v2.ONLY_FIRST_REPLICA, v1.ONLY_FIRST_TOWER)
+
+    self.assertEqual(hash(v1.NONE), hash(v2.NONE))
+    self.assertEqual(hash(v1.SUM), hash(v2.SUM))
+    self.assertEqual(hash(v1.MEAN), hash(v2.MEAN))
+    self.assertEqual(hash(v1.ONLY_FIRST_REPLICA), hash(v2.ONLY_FIRST_REPLICA))
+    self.assertEqual(hash(v1.ONLY_FIRST_TOWER), hash(v2.ONLY_FIRST_REPLICA))
+
+if __name__ == "__main__":
+  test.main()