def testRandomInitializer(self):
   # Sanity check that the slices uses a different seed when using a random
   # initializer function.
   with self.test_session():
     var0, var1 = partitioned_variables.create_partitioned_variables(
         [20, 12], [1, 2], init_ops.random_uniform_initializer())
     variables.global_variables_initializer().run()
     val0, val1 = var0.eval().flatten(), var1.eval().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.test_session():
     var0, var1 = partitioned_variables.create_partitioned_variables(
         [20, 12], [1, 2], init_ops.random_uniform_initializer(seed=201))
     variables.global_variables_initializer().run()
     val0, val1 = var0.eval().flatten(), var1.eval().flatten()
     self.assertAllClose(val0, val1)

 def testVecConstantInit(self):
   with self.test_session():
     rnd_par = constant_op.constant([1, 2, 3, 4])
     vs = partitioned_variables.create_partitioned_variables([4], [4], rnd_par)
     variables.global_variables_initializer().run()
     val = array_ops.concat(vs, 0).eval()
     rnd = rnd_par.eval()
     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 testDegenerate(self):
   with self.test_session():
     rnd = variables.Variable(random_ops.random_uniform([10, 43]))
     vs = partitioned_variables.create_partitioned_variables(
         rnd.get_shape(), [1, 1], rnd.initialized_value())
     variables.global_variables_initializer().run()
     val = array_ops.concat(vs, 0).eval()
     rnd = rnd.eval()
     self.assertAllClose(rnd, val)
     self._TestSaveSpec(vs, ["10 43 0,10:0,43"])

 def testConstantInit(self):
   with self.test_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)
     variables.global_variables_initializer().run()
     val = array_ops.concat(vs, 1).eval()
     rnd = rnd_par.eval()
     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 _random_weights(self, size=50, num_shards=1):
    assert size > 0
    assert num_shards > 0
    assert num_shards <= size

    embedding_weights = partitioned_variables.create_partitioned_variables(
        shape=[size],
        slicing=[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 _create_embedding_lookup(input_tensor, vocab_size, dimension,
                             weight_collections, initializer, combiner,
                             trainable, name):
  """Creates embedding variable and does a lookup.

  Args:
    input_tensor: A tensor which should contain sparse id to look up.
    vocab_size: An integer specifying the vocabulary size.
    dimension: An integer specifying the embedding vector dimension.
    weight_collections: List of graph collections to which weights are added.
    initializer: A variable initializer function to be used in embedding
      variable initialization.
    combiner: A string specifying how to reduce if the sparse column is
      multivalent. Currently "mean", "sqrtn" and "sum" are supported:
        * "sum": do not normalize features in the column
        * "mean": do l1 normalization on features in the column
        * "sqrtn": do l2 normalization on features in the column
      For more information: `tf.embedding_lookup_sparse`.
    trainable: If `True` also add variables to the graph collection
      `GraphKeys.TRAINABLE_VARIABLES` (see tf.Variable).
    name: A string specifying the name of the embedding variable.

  Returns:
    A Tensor with shape [batch_size, dimension] and embedding Variable.

  Raises:
    ValueError: If initializer is None or not callable.
  """
  slicing = _max_size_embedding_partitioner()(vocab_size, dimension)
  logging.info("Slicing=%s for name=%s, vocab_size=%d, embed_dim=%d",
               str(slicing), name, vocab_size, dimension)
  if not initializer:
    raise ValueError("initializer must be defined.")
  if not callable(initializer):
    raise ValueError("initializer must be callable.")
  embeddings = partitioned_variables.create_partitioned_variables(
      shape=[vocab_size, dimension],
      slicing=slicing,
      initializer=initializer,
      dtype=dtypes.float32,
      collections=weight_collections,
      name=name,
      reuse=False,
      trainable=trainable)

  return contrib_embedding_ops.safe_embedding_lookup_sparse(
      embeddings,
      input_tensor,
      default_id=0,
      combiner=combiner,
      name=name), embeddings

 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.test_session():
     vs = partitioned_variables.create_partitioned_variables([13, 5], [3, 1],
                                                             _IotaInitializer)
     variables.global_variables_initializer().run()
     slice0 = _IotaInitializer([5, 5])
     slice1 = _IotaInitializer([4, 5])
     slice2 = _IotaInitializer([4, 5])
     val = array_ops.concat(vs, 0).eval()
     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"])

 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], rnd.initialized_value())
     variables.global_variables_initializer().run()
     val = array_ops.concat(vs, 0).eval()
     rnd = rnd.eval()
     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 testRandomInitValue(self):
   with self.test_session():
     rnd = variables.Variable(random_ops.random_uniform([200, 40]))
     vs = partitioned_variables.create_partitioned_variables(
         rnd.get_shape(), [1, 10], rnd.initialized_value())
     variables.global_variables_initializer().run()
     val = array_ops.concat(vs, 1).eval()
     rnd = rnd.eval()
     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 _testNameHelper(self, use_resource=False):
   with self.test_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)
     variables.global_variables_initializer().run()
     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.test_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)
     variables.global_variables_initializer().run()
     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.test_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)
     variables.global_variables_initializer().run()
     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)

 def testLargePartitionedVariables(self):
   save_path = os.path.join(self.get_temp_dir(), "large_variable")
   var_name = "my_var"
   # Saving large partition variable.
   with session.Session("", graph=ops.Graph()) as sess:
     with ops.device("/cpu:0"):
       # Create a partitioned variable which is larger than int32 size but
       # split into smaller sized variables.
       init = lambda shape, dtype, partition_info: constant_op.constant(
           True, dtype, shape)
       partitioned_var = partitioned_variables.create_partitioned_variables(
           [1 << 31], [4], init, dtype=dtypes.bool, name=var_name)
       variables.global_variables_initializer().run()
       save = saver.Saver(partitioned_var)
       val = save.save(sess, save_path)
       self.assertEqual(save_path, val)

 def testRandomInitUnevenPartitions(self):
   with self.test_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], rnd.initialized_value())
         for i in xrange(1, 10)
     ]
     variables.global_variables_initializer().run()
     rnd_val = rnd.eval()
     # 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).eval()
       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 _create_partitioned_variables(name,
                                  num_hosts,
                                  vocabulary_size,
                                  embedding_dimension,
                                  initializer,
                                  collections=None):  # pylint: disable=redefined-outer-name
  """Creates ParitionedVariables based on `num_hosts` for `table`."""
  # TODO(shizhiw): automatically place embedding lookup elsewhere?
  if vocabulary_size < num_hosts:
    raise ValueError('`vocabulary_size`({}) is smaller than `num_hosts`({}). '
                     'As TPU embedding is not optimized for small tables, '
                     'please consider other ways for this embedding lookup.')

  slicing = [num_hosts, 1]

  # TODO(shizhiw): deprecated, use tf.get_variable()?
  return partitioned_variables.create_partitioned_variables(
      name=name,
      slicing=slicing,
      shape=(vocabulary_size, embedding_dimension),
      dtype=dtypes.float32,
      initializer=initializer,
      collections=collections,
      trainable=False)

 def testSomeErrors(self):
   with self.test_session():
     rnd = variables.Variable(random_ops.random_uniform([10, 43]))
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10], [1, 1], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 20], [1], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 43], [1], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 43], [1, 2, 3], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 43], [11, 1], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 43], [20, 1], rnd.initialized_value())
     with self.assertRaises(ValueError):
       partitioned_variables.create_partitioned_variables(
           [10, 43], [1, 50], rnd.initialized_value())