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

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

 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, variables.is_variable_initialized(v0).eval())
       state_ops.assign(v0, [[2.0, 3.0]]).eval()
       self.assertEqual(True, variables.is_variable_initialized(v0).eval())

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

 def testAssignNoShapeNoValidateShape(self):
   with self.test_session():
     value = self._NewShapelessTensor()
     var = state_ops.variable_op([1, 2], tf.float32, set_shape=False)
     self.assertEqual(tensor_shape.unknown_shape(), var.get_shape())
     self.assertEqual(tensor_shape.unknown_shape(),
                      tf.assign(var, value, validate_shape=False).get_shape())

  def __init__(self, initial_value, trainable=True, collections=None,
               validate_shape=True, name=None):
    """Creates a new variable with value `initial_value`.

    The new variable is added to the graph collections listed in `collections`,
    which defaults to `[GraphKeys.VARIABLES]`.

    If `trainable` is `True` the variable is also added to the graph collection
    `GraphKeys.TRAINABLE_VARIABLES`.

    This constructor creates both a `variable` Op and an `assign` Op to set the
    variable to its initial value.

    Args:
      initial_value: A `Tensor`, or Python object convertible to a `Tensor`.
        The initial value for the Variable. Must have a shape specified unless
        `validate_shape` is set to False.
      trainable: If `True`, the default, also adds the variable to the graph
        collection `GraphKeys.TRAINABLE_VARIABLES`. This collection is used as
        the default list of variables to use by the `Optimizer` classes.
      collections: List of graph collections keys. The new variable is added to
        these collections. Defaults to `[GraphKeys.VARIABLES]`.
      validate_shape: If `False`, allows the variable to be initialized with a
        value of unknown shape. If `True`, the default, the shape of
        `initial_value` must be known.
      name: Optional name for the variable. Defaults to `'Variable'` and gets
        uniquified automatically.

    Returns:
      A Variable.

    Raises:
      ValueError: If the initial value does not have a shape and
        `validate_shape` is `True`.
    """
    if collections is None:
      collections = [ops.GraphKeys.VARIABLES]
    if trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections:
      collections = list(collections) + [ops.GraphKeys.TRAINABLE_VARIABLES]
    with ops.control_dependencies(None):
      with ops.op_scope([initial_value], name, "Variable") as name:
        self._initial_value = ops.convert_to_tensor(initial_value,
                                                    name="initial_value")
        initial_value_shape = self._initial_value.get_shape()
        if validate_shape and not initial_value_shape.is_fully_defined():
          raise ValueError("initial_value must have a shape specified: %s"
                           % self._initial_value)
        shape_to_set = initial_value_shape if validate_shape else []
        self._variable = state_ops.variable_op(
            shape_to_set, self._initial_value.dtype.base_dtype,
            set_shape=validate_shape, name=name)
        with ops.device(self._variable.device):
          self._initializer_op = state_ops.assign(
              self._variable, self._initial_value,
              validate_shape=validate_shape).op
          self._snapshot = array_ops.identity(self._variable, name="read")

    ops.add_to_collections(collections, self)
    self._save_slice_info = None

  def _init_from_args(self, initial_value=None, trainable=True,
                      collections=None, validate_shape=True,
                      caching_device=None, name=None):
    """Creates a new variable from arguments.

    Args:
      initial_value: A `Tensor`, or Python object convertible to a `Tensor`.
        The initial value for the Variable. Must have a shape specified unless
        `validate_shape` is set to False.
      trainable: If `True`, the default, also adds the variable to the graph
        collection `GraphKeys.TRAINABLE_VARIABLES`. This collection is used as
        the default list of variables to use by the `Optimizer` classes.
      collections: List of graph collections keys. The new variable is added to
        these collections. Defaults to `[GraphKeys.VARIABLES]`.
      validate_shape: If `False`, allows the variable to be initialized with a
        value of unknown shape. If `True`, the default, the shape of
        `initial_value` must be known.
      caching_device: Optional device string or function describing where the
        Variable should be cached for reading.  Defaults to the Variable's
        device.  If not `None`, caches on another device.  Typical use is to
        cache on the device where the Ops using the Variable reside, to
        deduplicate copying through `Switch` and other conditional statements.
      name: Optional name for the variable. Defaults to `'Variable'` and gets
        uniquified automatically.

    Raises:
      ValueError: If the initial value is not specified, or does not have a
        shape and `validate_shape` is `True`.
    """
    if initial_value is None:
      raise ValueError("initial_value must be specified.")
    if collections is None:
      collections = [ops.GraphKeys.VARIABLES]
    if trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections:
      collections = list(collections) + [ops.GraphKeys.TRAINABLE_VARIABLES]
    with ops.control_dependencies(None):
      with ops.op_scope([initial_value], name, "Variable") as name:
        self._initial_value = ops.convert_to_tensor(initial_value,
                                                    name="initial_value")
        initial_value_shape = self._initial_value.get_shape()
        if validate_shape and not initial_value_shape.is_fully_defined():
          raise ValueError("initial_value must have a shape specified: %s"
                           % self._initial_value)
        shape_to_set = initial_value_shape if validate_shape else []
        self._variable = state_ops.variable_op(
            shape_to_set, self._initial_value.dtype.base_dtype,
            set_shape=validate_shape, name=name)
        with ops.device(self._variable.device):
          self._initializer_op = state_ops.assign(
              self._variable, self._initial_value,
              validate_shape=validate_shape).op
        with ops.device(caching_device if caching_device is not None
                        else self._variable.device):
          self._snapshot = array_ops.identity(self._variable, name="read")

    ops.add_to_collections(collections, self)
    self._caching_device = caching_device
    self._save_slice_info = None

 def testAssignDependencyAcrossDevices(self):
   with self.test_session(use_gpu=True):
     # The variable and an op to increment it are on the GPU.
     var = state_ops.variable_op([1], tf.float32)
     tf.assign(var, [1.0]).eval()
     increment = tf.assign_add(var, [1.0])
     with tf.control_dependencies([increment]):
       with tf.device("/cpu:0"):
         # This mul op is pinned to the CPU, but reads the variable from the
         # GPU. The te

 def testAverageVariablesDeviceAssignment(self):
   with ops.device("dev_v0"):
     v0 = variables.Variable(10.0, name="v0")
   with ops.device("dev_v1"):
     v1 = state_ops.variable_op(shape=[1], dtype=types.float32, name="v1")
   tensor2 = v0 + v1
   ema = moving_averages.ExponentialMovingAverage(0.25, name="foo_avg")
   with ops.device("default"):
     ema.apply([v0, v1, tensor2])
   self.assertEqual("dev_v0", ema.average(v0).device)
   self.assertEqual("dev_v1", ema.average(v1).device)
   self.assertEqual("default", ema.average(tensor2).device)

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

 def _buildInitialVars(self, shape, dev_list):
   values = []
   num_devices = len(dev_list)
   dim = np.prod(shape) if shape else 1
   for d in range(0, num_devices):
     with ops.device(dev_list[d]):
       npt = np.zeros(shape).astype(np.float32)
       alias = np.frombuffer(npt.data, dtype=np.float32)
       for i in range(0, dim):
         alias[i] = i + 0.01 * d
       var = state_ops.variable_op(shape, types_pb2.DT_FLOAT)
       state_ops.init_variable(var, npt).op.run()
       values.append(var)
   return values

 def testObtainNext(self):
   with self.test_session():
     var = state_ops.variable_op([1], tf.int64)
     tf.assign(var, [-1]).op.run()
     c = tf.constant(["a", "b"])
     sample1 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"a", sample1.eval())
     self.assertEqual([0], var.eval())
     sample2 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"b", sample2.eval())
     self.assertEqual([1], var.eval())
     sample3 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"a", sample3.eval())
     self.assertEqual([0], var.eval())

 def testObtainNext(self):
   with self.test_session():
     var = state_ops.variable_op([], dtypes.int64)
     state_ops.assign(var, -1).op.run()
     c = constant_op.constant(["a", "b"])
     sample1 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"a", sample1.eval())
     self.assertEqual(0, var.eval())
     sample2 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"b", sample2.eval())
     self.assertEqual(1, var.eval())
     sample3 = input_pipeline_ops.obtain_next(c, var)
     self.assertEqual(b"a", sample3.eval())
     self.assertEqual(0, var.eval())

 def testAssignDependencyAcrossDevices(self):
   with self.test_session(use_gpu=True):
     # The variable and an op to increment it are on the GPU.
     var = state_ops.variable_op([1], tf.float32)
     tf.assign(var, [1.0]).eval()
     increment = tf.assign_add(var, [1.0])
     with tf.control_dependencies([increment]):
       with tf.device("/cpu:0"):
         # 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 = tf.mul(var, var)
     self.assertAllClose([4.0], result.eval())

 def testAverageVariablesDeviceAssignment(self):
   with tf.device("/job:dev_v0"):
     v0 = tf.Variable(10.0, name="v0")
   with tf.device("/job:dev_v1"):
     v1 = state_ops.variable_op(shape=[1], dtype=tf.float32, name="v1")
   tensor2 = v0 + v1
   ema = tf.train.ExponentialMovingAverage(0.25, name="foo_avg")
   with tf.device("/job:default"):
     ema.apply([v0, v1, tensor2])
   self.assertDeviceEqual("/job:dev_v0", ema.average(v0).device)
   self.assertDeviceEqual("/job:dev_v1", ema.average(v1).device)
   # However, the colocation property is maintained.
   self.assertEqual([b"loc:@v1"],
                    ema.average(v1).op.colocation_groups())
   self.assertDeviceEqual("/job:default", ema.average(tensor2).device)

 def testDecay(self):
   initial_lr = 0.1
   k = 10
   decay_rate = 0.96
   step = state_ops.variable_op([], dtypes.int32)
   assign_step = state_ops.assign(step, 0)
   increment_step = state_ops.assign_add(step, 1)
   decayed_lr = learning_rate_decay.natural_exp_decay(initial_lr, step,
                                                      k, decay_rate)
   with self.test_session():
     assign_step.op.run()
     for i in range(k+1):
       expected = initial_lr * math.exp(-i / k * decay_rate)
       self.assertAllClose(decayed_lr.eval(), expected, 1e-6)
       increment_step.op.run()

 def testStaircase(self):
     with self.test_session():
         step = state_ops.variable_op([], dtypes.int32)
         assign_100 = state_ops.assign(step, 100)
         assign_1 = state_ops.assign(step, 1)
         assign_2 = state_ops.assign(step, 2)
         decayed_lr = learning_rate_decay.exponential_decay(0.1, step, 3, 0.96, staircase=True)
         # No change to learning rate
         assign_1.op.run()
         self.assertAllClose(decayed_lr.eval(), 0.1, 1e-6)
         assign_2.op.run()
         self.assertAllClose(decayed_lr.eval(), 0.1, 1e-6)
         # Decayed learning rate
         assign_100.op.run()
         expected = 0.1 * 0.96 ** (100 // 3)
         self.assertAllClose(decayed_lr.eval(), expected, 1e-6)

 def testContainer(self):
   with tf.Graph().as_default():
     v0 = tf.Variable([0])
     with tf.container("l1"):
       v1 = tf.Variable([1])
       with tf.container("l2"):
         v2 = tf.Variable([2])
         special_v = state_ops.variable_op([1], tf.float32, container="l3")
       v3 = tf.Variable([3])
     v4 = tf.Variable([4])
   self.assertEqual(tf.compat.as_bytes(""), v0.op.get_attr("container"))
   self.assertEqual(tf.compat.as_bytes("l1"), v1.op.get_attr("container"))
   self.assertEqual(tf.compat.as_bytes("l2"), v2.op.get_attr("container"))
   self.assertEqual(tf.compat.as_bytes("l3"),
                    special_v.op.get_attr("container"))
   self.assertEqual(tf.compat.as_bytes("l1"), v3.op.get_attr("container"))
   self.assertEqual(tf.compat.as_bytes(""), v4.op.get_attr("container"))

 def testStaircase(self):
   initial_lr = 0.1
   k = 10
   decay_rate = 0.96
   step = state_ops.variable_op([], dtypes.int32)
   assign_step = state_ops.assign(step, 0)
   increment_step = state_ops.assign_add(step, 1)
   decayed_lr = learning_rate_decay.inverse_time_decay(initial_lr,
                                                       step,
                                                       k,
                                                       decay_rate,
                                                       staircase=True)
   with self.test_session():
     assign_step.op.run()
     for i in range(k+1):
       expected = initial_lr / (1 + decay_rate * (i // k))
       self.assertAllClose(decayed_lr.eval(), expected, 1e-6)
       increment_step.op.run()

 def testPinRequiredOpsOnCPU(self):
     with ops.Graph().as_default() as g, g.device(graph_util.pin_variables_on_cpu):
         const_a = constant_op.constant(5.0)
         const_b = constant_op.constant(10.0)
         add_c = const_a + const_b
         var_v = state_ops.variable_op([], dtype=types.float32)
         assign_c_to_v = state_ops.assign(var_v, add_c)
         dynamic_stitch_int_result = data_flow_ops.dynamic_stitch([[0, 1, 2], [2, 3]], [[12, 23, 34], [1, 2]])
         dynamic_stitch_float_result = data_flow_ops.dynamic_stitch(
             [[0, 1, 2], [2, 3]], [[12.0, 23.0, 34.0], [1.0, 2.0]]
         )
         # Non-variable ops shuld not specify a device
         self.assertEqual(const_a.device, None)
         self.assertEqual(const_b.device, None)
         self.assertEqual(add_c.device, None)
         # Variable ops specify a device
         self.assertEqual(var_v.device, "/device:CPU:0")
         self.assertEqual(assign_c_to_v.device, "/device:CPU:0")