def _GenerateTestInputs(self):
    np.random.seed(0)
    weights = np.random.randn(self._num_classes, self._dim).astype(np.float32)
    biases = np.random.randn(self._num_classes).astype(np.float32)
    hidden_acts = np.random.randn(self._batch_size,
                                  self._dim).astype(np.float32)

    with ops.Graph().as_default() as g:
      sharded_weights = variable_scope.get_variable(
          "w",
          partitioner=partitioned_variables.fixed_size_partitioner(
              self._num_shards),
          initializer=constant_op.constant(weights))
      sharded_biases = variable_scope.get_variable(
          "b",
          partitioner=partitioned_variables.fixed_size_partitioner(
              self._num_shards),
          initializer=constant_op.constant(biases))
      with self.test_session(graph=g) as sess:
        variables.global_variables_initializer().run()

        sharded_weights_v, sharded_biases_v = sess.run(
            [list(sharded_weights), list(sharded_biases)])

    return weights, biases, hidden_acts, sharded_weights_v, sharded_biases_v

  def test_load_and_remap_linear_multiclass_initializer_default_init(self):
    """Tests where the zeros_initializer default is used for linear."""
    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
        new_row_vocab_size=5,
        new_col_vocab_file=self.new_class_vocab_file,
        old_col_vocab_file=self.old_class_vocab_file,
        new_col_vocab_size=4,
        old_tensor_name='some_scope/embeddings',
        ckpt_path=[self.checkpoint_file],
        new_row_vocab_file=self.new_feature_vocab_file,
        old_row_vocab_file=self.old_feature_vocab_file,
        num_row_oov_buckets=1,
        num_col_oov_buckets=1))

    expected_remapped_matrix = np.concatenate(
        [
            np.reshape([2, 18, 34, 50, 0, 0], [6, 1]),
            np.reshape([0, 16, 32, 48, 0, 0], [6, 1]),
            np.reshape([0] * 6, [6, 1]),
            np.reshape([1, 17, 33, 49, 0, 0], [6, 1]),
            np.reshape([0] * 6, [6, 1])
        ],
        axis=1)

    remapped_matrix = variable_scope.get_variable(
        name='linear_init_fallback/obtained_weight_matrix',
        shape=[6, 5],
        initializer=loading_initializer,
        partitioner=partitioned_variables.fixed_size_partitioner(2))

    with self.test_session():
      variables.global_variables_initializer().run()
      self.assertAllClose(expected_remapped_matrix,
                          remapped_matrix.as_tensor().eval())

  def test_load_and_remap_output_layer_weight_initializer_dnn_output(self):
    """Tests for the output layer initializer in the DNN output case."""
    loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer(
        new_row_vocab_size=5,
        new_col_vocab_file=self.new_class_vocab_file,
        old_col_vocab_file=self.old_class_vocab_file,
        new_col_vocab_size=4,
        old_tensor_name='some_scope/embeddings',
        ckpt_path=[self.checkpoint_file],
        num_col_oov_buckets=1,
        initializer=self.initializer))

    expected_remapped_matrix = np.concatenate(
        [
            np.reshape([2, 18, 34, 50, 66], [5, 1]),
            np.reshape([0, 16, 32, 48, 64], [5, 1]),
            np.reshape([self.init_val] * 5, [5, 1]),
            np.reshape([1, 17, 33, 49, 65], [5, 1]),
            np.reshape([self.init_val] * 5, [5, 1])
        ],
        axis=1)

    # The new weight matrix is of size
    # [5-sized input layer, 4 class vocab + 1 class OOV].
    remapped_matrix = variable_scope.get_variable(
        name='dnn_output/obtained_weight_matrix',
        shape=[5, 5],
        initializer=loading_initializer,
        partitioner=partitioned_variables.fixed_size_partitioner(2))

    with self.test_session():
      variables.global_variables_initializer().run()
      self.assertAllClose(expected_remapped_matrix,
                          remapped_matrix.as_tensor().eval())

 def testFixedSizePartitionerInt64(self):
   with self.test_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 test_load_linear_multiclass_bias_initializer(self):
    """Tests for the bias initializer wrapper."""
    bias_loading_initializer = (
        contrib_framework.load_linear_multiclass_bias_initializer(
            new_class_vocab_file=self.new_class_vocab_file,
            old_class_vocab_file=self.old_class_vocab_file,
            new_class_vocab_size=4,
            bias_tensor_name='some_scope/bias',
            ckpt_path=[self.bundle_file],
            num_class_oov_buckets=1,
            initializer=self.initializer))

    expected_remapped_bias_vector = np.reshape(
        [2, 0, self.init_val, 1, self.init_val], [5, 1])

    # The new bias vector is of size [4 class vocab + 1 class OOV, 1].
    remapped_bias_vector = variable_scope.get_variable(
        name='bias/obtained_bias_vector',
        shape=[5, 1],
        initializer=bias_loading_initializer,
        partitioner=partitioned_variables.fixed_size_partitioner(3))

    with self.test_session():
      variables.global_variables_initializer().run()
      self.assertAllClose(expected_remapped_bias_vector,
                          remapped_bias_vector.as_tensor().eval())

  def _testMultiplePartitionedVariables(self, is_training):
    # When weights are partitioned into multiple partitions the weights variable
    # is followed by a identity -> concat -> identity to group the partitions.
    partitioner = partitioned_variables.fixed_size_partitioner(2)
    graph = ops.Graph()
    with graph.as_default():
      with variable_scope.variable_scope('part', partitioner=partitioner):
        batch_size, height, width, depth = 5, 128, 128, 3
        input1 = array_ops.zeros((batch_size, height, width, depth))
        input2 = array_ops.zeros((batch_size, height / 2, width / 2, 32))
        conv = conv2d(
            input1,
            32, [5, 5],
            stride=2,
            padding='SAME',
            weights_initializer=self._WeightInit(0.09),
            activation_fn=None,
            scope='test/test')
        node = math_ops.add(conv, input2, name='test/add')
        node = nn_ops.relu6(node, name='test/relu6')

      quantize.Quantize(graph, is_training, weight_bits=8, activation_bits=8)
      # Check that the weight's quant node was added.
      op_names = [op.name for op in graph.get_operations()]
      self.assertTrue(
          'part/test/test/weights_quant/FakeQuantWithMinMaxVars' in op_names)

  def testMetaGraphSaveLoad(self):
    save_prefix = os.path.join(self.get_temp_dir(), "ckpt")
    save_graph = ops.Graph()
    with save_graph.as_default(), self.test_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))
        variables.global_variables_initializer().run()

        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.test_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")))

  def test_load_embedding_initializer(self):
    """Tests for the load_embedding_initializer wrapper."""
    embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
        new_vocab_file=self.new_feature_vocab_file,
        old_vocab_file=self.old_feature_vocab_file,
        new_vocab_size=5,
        embedding_dim=16,
        embedding_tensor_name='some_scope/embeddings',
        ckpt_path=[self.checkpoint_file],
        num_oov_buckets=1,
        initializer=self.initializer))

    expected_remapped_embeddings = np.concatenate(
        [
            np.reshape(range(64), [4, 16]),
            np.reshape([self.init_val] * 32, [2, 16]),
        ],
        axis=0)

    # The new weight matrix is of size
    # [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the
    # last vocab row (2nd last row) is newly initialized (wasn't found in
    # previous vocab) and the actual last row is OOV and also newly initialized.
    # Use a partitioned variable to confirm that the offset logic works.
    remapped_embeddings = variable_scope.get_variable(
        name='embedding/obtained_embedding_matrix',
        shape=[6, 16],
        initializer=embedding_loading_initializer,
        partitioner=partitioned_variables.fixed_size_partitioner(2))

    with self.test_session():
      variables.global_variables_initializer().run()
      self.assertAllClose(expected_remapped_embeddings,
                          remapped_embeddings.as_tensor().eval())

 def test_loading_partitions_greater_than_max_rows(self):
   """Tests loading partitioned var with more slices than partitions."""
   self._test_loading_variable_with_max_rows(
       np_value=np.reshape(list(range(0, 36)), (9, 4)),
       partitioner=partitioned_variables.fixed_size_partitioner(3),
       # Even though each partition has 3 rows, we'll only load the tensor one
       # row at a time.
       max_rows_in_memory=1)

 def test_loading_partitions_equals_max_rows(self):
   """Tests loading partitioned var sliced on partition boundary."""
   self._test_loading_variable_with_max_rows(
       np_value=np.reshape(list(range(0, 36)), (9, 4)),
       partitioner=partitioned_variables.fixed_size_partitioner(3),
       # With a tensor of shape [9, 3] and 3 partitions, each partition has
       # exactly 3 rows.
       max_rows_in_memory=3)

  def test_loading_partitions_less_than_max_rows(self):
    """Tests loading partitioned var as a single slice.

    (When the specified max_rows_in_memory is larger than the number of rows)
    """
    self._test_loading_variable_with_max_rows(
        np_value=np.reshape(list(range(0, 36)), (9, 4)),
        partitioner=partitioned_variables.fixed_size_partitioner(3),
        max_rows_in_memory=10)

 def testFixedSizePartitioner(self):
   with self.test_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 build(self, inputs_shape):
    if inputs_shape[1].value is None:
      raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s"
                       % inputs_shape)

    input_depth = inputs_shape[1].value
    h_depth = self._num_units if self._num_proj is None else self._num_proj
    maybe_partitioner = (
        partitioned_variables.fixed_size_partitioner(self._num_unit_shards)
        if self._num_unit_shards is not None
        else None)
    self._kernel = self.add_variable(
        _WEIGHTS_VARIABLE_NAME,
        shape=[input_depth + h_depth, 4 * self._num_units],
        initializer=self._initializer,
        partitioner=maybe_partitioner)
    self._bias = self.add_variable(
        _BIAS_VARIABLE_NAME,
        shape=[4 * self._num_units],
        initializer=init_ops.zeros_initializer(dtype=self.dtype))
    if self._use_peepholes:
      self._w_f_diag = self.add_variable("w_f_diag", shape=[self._num_units],
                                         initializer=self._initializer)
      self._w_i_diag = self.add_variable("w_i_diag", shape=[self._num_units],
                                         initializer=self._initializer)
      self._w_o_diag = self.add_variable("w_o_diag", shape=[self._num_units],
                                         initializer=self._initializer)

    if self._num_proj is not None:
      maybe_proj_partitioner = (
          partitioned_variables.fixed_size_partitioner(self._num_proj_shards)
          if self._num_proj_shards is not None
          else None)
      self._proj_kernel = self.add_variable(
          "projection/%s" % _WEIGHTS_VARIABLE_NAME,
          shape=[self._num_units, self._num_proj],
          initializer=self._initializer,
          partitioner=maybe_proj_partitioner)

    self.built = True

 def make_graph_with_partitioned_variables(use_resource):
   variable_scope.get_variable(
       name="weights",
       partitioner=partitioned_variables.fixed_size_partitioner(3, axis=0),
       initializer=random_ops.truncated_normal([100, 10]),
       use_resource=use_resource)
   # The next variable illustrates the necessity of restoring collections
   # in a deterministic fashion when using ResourceVariables.
   variable_scope.get_variable(
       name="another",
       shape=[],
       collections=["a", "b", "z", "f", "e", "d", "g"],
       use_resource=use_resource)

 def testVariablesNotParitioned_MovingAvg(self):
   # Variables added should not use a default partiioner since they are
   # scalar. There would be a tensorflow error thrown if the partitioner was
   # respected by the rewrite.
   with ops.Graph().as_default():
     with variable_scope.variable_scope(
         'part', partitioner=partitioned_variables.fixed_size_partitioner(2)):
       x = array_ops.placeholder(dtypes.float32, shape=[2])
       _ = quant_ops.MovingAvgQuantize(
           x,
           init_min=0.0,
           init_max=0.0,
           is_training=True,
           vars_collection=_MIN_MAX_VARS)

  def _random_weights(self, size=50, num_shards=1):
    assert size > 0
    assert num_shards > 0
    assert num_shards <= size

    embedding_weights = list(variable_scope.get_variable(
        "embedding_weights",
        shape=[size],
        partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
        initializer=init_ops.truncated_normal_initializer(
            mean=0.0, stddev=1.0, dtype=dtypes.float32)))
    for w in embedding_weights:
      w.initializer.run()
    return embedding_weights

 def testReusePartitionedVaraiblesAndRegularizers(self):
   regularizer = lambda x: math_ops.reduce_sum(x) * 1e-3
   partitioner = partitioned_variables.fixed_size_partitioner(3)
   for reuse in [False, True]:
     with variable_scope.variable_scope(variable_scope.get_variable_scope(),
                                        partitioner=partitioner,
                                        reuse=reuse):
       layer = base_layers.Layer(name='my_layer')
       variable = layer.add_variable(
           'reg_part_var', [4, 4],
           initializer=init_ops.zeros_initializer(),
           regularizer=regularizer)
   self.assertEqual(
       len(ops.get_collection(ops.GraphKeys.REGULARIZATION_LOSSES)), 3)

  def _ShardTestEmbeddings(self, weights, biases, num_shards):
    """Shards the weights and biases returned by _GenerateTestData.

    Args:
      weights: The weights returned by _GenerateTestData.
      biases: The biases returned by _GenerateTestData.
      num_shards: The number of shards to create.

    Returns:
      sharded_weights: A list of size `num_shards` containing all the weights.
      sharded_biases: A list of size `num_shards` containing all the biases.
    """
    with ops.Graph().as_default() as g:
      sharded_weights = variable_scope.get_variable(
          "w",
          partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
          initializer=constant_op.constant(weights))
      sharded_biases = variable_scope.get_variable(
          "b",
          partitioner=partitioned_variables.fixed_size_partitioner(num_shards),
          initializer=constant_op.constant(biases))
      with self.test_session(graph=g) as sess:
        variables.global_variables_initializer().run()
        return sess.run([list(sharded_weights), list(sharded_biases)])

  def _test_device_assignment_distributed_enable_partitioner(
      self, task_type, task_id, num_gpus, use_core_strategy=False):
    d, _, sess_config = self._get_test_objects(
        task_type, task_id, num_gpus, use_core_strategy=use_core_strategy)
    num_shards = len(d.extended.parameter_devices)
    partitioner = partitioned_variables.fixed_size_partitioner(num_shards)
    with ops.Graph().as_default(), \
         self.cached_session(target=self._default_target,
                             config=sess_config) as sess, \
         d.scope():

      n = variable_scope.get_variable(
          'n',
          initializer=constant_op.constant([10.0, 20.0]),
          aggregation=variable_scope.VariableAggregation.SUM,
          partitioner=partitioner)

      for part_id, var in enumerate(n):
        self.assertEqual(var.device, '/job:ps/task:%d' % part_id)

      def model_fn():
        a = constant_op.constant([3.0, 5.0])
        # The device scope is ignored for variables but not for normal ops.
        with ops.device('/job:worker/task:0'):
          x = variable_scope.get_variable(
              'x',
              initializer=constant_op.constant([10.0, 20.0]),
              aggregation=variable_scope.VariableAggregation.SUM,
              partitioner=partitioner)
          x_add = x.assign_add(a, name='x_add')
        # The variable x is on the task 1 since the device_function has been
        # called once before the model_fn.
        for part_id, var in enumerate(x):
          self.assertEqual(var.device, '/job:ps/task:%d' % part_id)
          self.assertEqual(var.device, x_add[part_id].device)

        return x_add

      x = d.extended.call_for_each_replica(model_fn)

      if context.num_gpus() >= 1:
        variables.global_variables_initializer().run()
        x_val = sess.run(x)
        if num_gpus < 1:
          self.assertEqual(x_val, [13.0, 25.0])
        else:
          x_expect = [10.0 + 3 * num_gpus, 20.0 + 5 * num_gpus]
          self.assertEqual(x_val, x_expect)

  def test_load_embedding_initializer_large_oov(self):
    """Tests for the large OOV case for load_embedding_initializer wrapper."""
    self.new_feature_vocab_file = os.path.join(
        self.get_temp_dir(), 'new_feature_vocab.txt')
    with open(self.new_feature_vocab_file, 'w') as f:
      f.write('\n'.join(['one', 'zero', 'two', 'four']) + '\n')

    # Checkpoint has 5 entries, 3 of which correspond to OOV.
    self.old_feature_vocab_file = os.path.join(
        self.get_temp_dir(), 'old_feature_vocab.txt')
    with open(self.old_feature_vocab_file, 'w') as f:
      f.write('\n'.join(['zero', 'one']) + '\n')

    embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer(
        new_vocab_file=self.new_feature_vocab_file,
        old_vocab_file=self.old_feature_vocab_file,
        new_vocab_size=4,
        embedding_dim=16,
        embedding_tensor_name='some_scope/embeddings',
        ckpt_path=[self.checkpoint_file],
        num_oov_buckets=5,
        initializer=self.initializer))

    expected_remapped_embeddings = np.concatenate(
        [
            np.reshape(range(16, 32), [1, 16]),
            np.reshape(range(16), [1, 16]),
            np.reshape([self.init_val] * 112, [7, 16]),
        ],
        axis=0)

    # The new weight matrix is of size
    # [4 feature vocab + 5 feature OOV, 16 (embedding dimension)], where the
    # 3rd and 4th rows are not found in the old vocabulary and therefore newly
    # initialized.  The last five rows are OOV and also newly initialized.
    # Use a partitioned variable to confirm that the offset logic works.
    remapped_embeddings = variable_scope.get_variable(
        name='embedding/obtained_embedding_matrix',
        shape=[9, 16],
        initializer=embedding_loading_initializer,
        partitioner=partitioned_variables.fixed_size_partitioner(2))

    with self.test_session():
      variables.global_variables_initializer().run()
      self.assertAllClose(expected_remapped_embeddings,
                          remapped_embeddings.as_tensor().eval())