def testInitializerDifferent(self):
   for dtype in [dtypes.float32, dtypes.float64]:
     init1 = init_ops.truncated_normal_initializer(
         0.0, 1.0, seed=1, dtype=dtype)
     init2 = init_ops.truncated_normal_initializer(
         0.0, 1.0, seed=2, dtype=dtype)
     self.assertFalse(identicaltest(self, init1, init2))

  def testInitFromPartitionVar(self):
    checkpoint_dir = self.get_temp_dir()
    with self.test_session() as session:
      v1 = _create_partition_checkpoints(session, checkpoint_dir)

    # New graph and session.
    with ops.Graph().as_default() as g:
      with self.test_session(graph=g) as session:
        with variable_scope.variable_scope("some_scope"):
          my1 = variable_scope.get_variable(
              name="my1",
              shape=[100, 100],
              initializer=init_ops.truncated_normal_initializer(0.5),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=8 << 10))
          my1_var_list = my1._get_variable_list()
        with variable_scope.variable_scope("some_other_scope"):
          my2 = variable_scope.get_variable(
              name="var1",
              shape=[100, 100],
              initializer=init_ops.truncated_normal_initializer(0.5),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=8 << 10))
          my2_var_list = my2._get_variable_list()

        checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
            "scope/var1": "some_scope/my1",
            "scope/": "some_other_scope/"})

        session.run(variables.global_variables_initializer())
        my1_values = session.run(my1_var_list)
        self.assertAllEqual(my1_values, v1)
        my2_values = session.run(my2_var_list)
        self.assertAllEqual(my2_values, v1)

    # New graph and session.
    with ops.Graph().as_default() as g:
      with self.test_session(graph=g) as session:
        with variable_scope.variable_scope("some_scope"):
          my1 = variable_scope.get_variable(
              name="my1",
              shape=[100, 100],
              initializer=init_ops.truncated_normal_initializer(0.5),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=8 << 10))
          my1_var_list = my1._get_variable_list()

        checkpoint_utils.init_from_checkpoint(checkpoint_dir,
                                              {"scope/var1": my1_var_list,})

        session.run(variables.global_variables_initializer())
        my1_values = session.run(my1_var_list)
        self.assertAllEqual(my1_values, v1)

  def _define_vars(self, params, **kwargs):
    with ops.device(self.device_assigner):

      self.tree_parameters = variable_scope.get_variable(
          name='tree_parameters_%d' % self.layer_num,
          shape=[params.num_nodes, params.num_features_per_node],
          initializer=init_ops.truncated_normal_initializer(
              mean=params.weight_init_mean, stddev=params.weight_init_std))

      self.tree_thresholds = variable_scope.get_variable(
          name='tree_thresholds_%d' % self.layer_num,
          shape=[params.num_nodes],
          initializer=init_ops.truncated_normal_initializer(
              mean=params.weight_init_mean, stddev=params.weight_init_std))

  def testWithScopes(self):
    init_value0 = np.asarray([1.0, 3.0, 9.0]).reshape((1, 3, 1))
    init_value1 = np.asarray([2.0, 4.0, 6.0, 8.0]).reshape((2, 1, 2))

    with self.test_session() as sess:
      initializer = init_ops.truncated_normal_initializer(stddev=.1)

      with variable_scope.variable_scope('my_model/my_layer0'):
        var0 = variables_lib2.variable(
            'my_var0', shape=[1, 3, 1], initializer=initializer)
      with variable_scope.variable_scope('my_model/my_layer1'):
        var1 = variables_lib2.variable(
            'my_var1', shape=[2, 1, 2], initializer=initializer)

      var_names_to_values = {
          'my_model/my_layer0/my_var0': init_value0,
          'my_model/my_layer1/my_var1': init_value1
      }
      init_fn = variables_lib2.assign_from_values_fn(var_names_to_values)

      # Initialize the variables.
      sess.run(variables_lib.global_variables_initializer())

      # Perform the assignment.
      init_fn(sess)

      # Request and test the variable values:
      var0, var1 = sess.run([var0, var1])
      self.assertAllEqual(init_value0, var0)
      self.assertAllEqual(init_value1, var1)

def load_embedding_initializer(ckpt_path,
                               embedding_tensor_name,
                               new_vocab_size,
                               embedding_dim,
                               old_vocab_file,
                               new_vocab_file,
                               num_oov_buckets=0,
                               initializer=None):
  """Returns a variable initializer for loading pre-trained embeddings.

  Wrapper around `load_and_remap_matrix_initializer()` specialized for loading
  embedding weights and remapping according to the provided vocab files. See
  docs for `load_and_remap_matrix_initializer()` for more details.

  NOTE: Only for use with div-partitioned variables / vocabularies.

  Args:
    ckpt_path: Path to the TensorFlow checkpoint (version 2, `TensorBundle`)
      from which the old matrix `Tensor` will be loaded.
    embedding_tensor_name: Name of the 2-D `Tensor` to load from checkpoint.
    new_vocab_size: Number of entries in the new vocab.
    embedding_dim: `int` specifying the dimension of the embedding vectors from
      the checkpoint. Must match the number of columns in the old embedding
      matrix.
    old_vocab_file: A scalar `Tensor` of type `string` containing the
      path to the old vocabulary file.
    new_vocab_file: A scalar `Tensor` of type `string` containing the
      path to the new vocabulary file.
    num_oov_buckets: `int` specifying the number of out-of-vocabulary
      buckets to use. Must be >= 0.
    initializer: Initializer function that accepts a 1-D tensor as the arg to
      specify the shape of the returned tensor. If `None`, defaults to using
      `truncated_normal_initializer()`.

  Returns:
    A variable initializer function.
  """
  if initializer is None:
    # TODO(b/25671353): This should be kept in sync with the stddev used by
    # feature_column.py's _EmbeddingColumn.
    initializer = init_ops.truncated_normal_initializer(
        stddev=1.0 /
        math_ops.sqrt(math_ops.cast(embedding_dim, dtypes.float32)))

  return load_and_remap_matrix_initializer(
      ckpt_path=ckpt_path,
      old_tensor_name=embedding_tensor_name,
      new_row_vocab_size=new_vocab_size,
      new_col_vocab_size=embedding_dim,
      old_row_vocab_file=old_vocab_file,
      new_row_vocab_file=new_vocab_file,
      old_col_vocab_file=None,
      new_col_vocab_file=None,
      num_row_oov_buckets=num_oov_buckets,
      num_col_oov_buckets=0,
      initializer=initializer)

  def _WeightInit(self, stddev):
    """Returns truncated normal variable initializer.

    Function is defined purely to shorten the name so that it stops wrapping.

    Args:
      stddev: Standard deviation of normal variable.

    Returns:
      An initialized that initializes with a truncated normal variable.
    """
    return init_ops.truncated_normal_initializer(stddev=stddev)

  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 _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 __new__(cls,
             sparse_id_column,
             dimension,
             combiner="mean",
             initializer=None):
   if initializer is not None and not callable(initializer):
     raise ValueError("initializer must be callable if specified.")
   if initializer is None:
     stddev = 1 / math.sqrt(sparse_id_column.length)
     # TODO(b/25671353): Better initial value?
     initializer = init_ops.truncated_normal_initializer(mean=0.0,
                                                         stddev=stddev)
   return super(_EmbeddingColumn, cls).__new__(cls, sparse_id_column,
                                               dimension, combiner,
                                               initializer)

  def __new__(cls,
              vocabulary_size,
              dimension,
              initializer=None,
              combiner='mean'):
    """Embedding table configuration.

    Args:
      vocabulary_size: Number of vocabulary (/rows) in the table.
      dimension: The embedding dimension.
      initializer: A variable initializer function to be used in embedding
        variable initialization. If not specified, defaults to
        `tf.truncated_normal_initializer` with mean `0.0` and standard deviation
        `1/sqrt(dimension)`.
      combiner: A string specifying how to reduce if there are multiple entries
        in a single row. Currently 'mean', 'sqrtn', 'sum' and None are
        supported, with 'mean' the default. 'sqrtn' often achieves good
        accuracy, in particular with bag-of-words columns. For more information,
        see `tf.nn.embedding_lookup_sparse`. None is only valid for dense rather
        than sparse tensors.

    Returns:
      `TableConfig`.

    Raises:
      ValueError: if `vocabulary_size` is not positive integer.
      ValueError: if `dimension` is not positive integer.
      ValueError: if `initializer` is specified and is not callable.
      ValueError: if `combiner` is not supported.
    """
    if not isinstance(vocabulary_size, int) or vocabulary_size < 1:
      raise ValueError('Invalid vocabulary_size {}.'.format(vocabulary_size))

    if not isinstance(dimension, int) or dimension < 1:
      raise ValueError('Invalid dimension {}.'.format(dimension))

    if (initializer is not None) and (not callable(initializer)):
      raise ValueError('initializer must be callable if specified.')
    if initializer is None:
      initializer = init_ops.truncated_normal_initializer(
          mean=0.0, stddev=1 / math.sqrt(dimension))

    if combiner not in ('mean', 'sum', 'sqrtn', None):
      raise ValueError('Invalid combiner {}'.format(combiner))

    return super(TableConfig, cls).__new__(cls, vocabulary_size, dimension,
                                           initializer, combiner)

def _create_embedding_lookup(input_tensor, vocab_size, dimension,
                             weight_collections, stddev, 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.
    stddev: the standard deviation to be used in embedding 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.
  """
  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 stddev > 0:
    initializer = init_ops.truncated_normal_initializer(stddev=stddev)
  else:
    initializer = init_ops.zeros_initializer
  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 inception_v3_arg_scope(weight_decay=0.00004,
                           stddev=0.1,
                           batch_norm_var_collection='moving_vars',
                           use_fused_batchnorm=True):
  """Defines the default InceptionV3 arg scope.

  Args:
    weight_decay: The weight decay to use for regularizing the model.
    stddev: The standard deviation of the trunctated normal weight initializer.
    batch_norm_var_collection: The name of the collection for the batch norm
      variables.
    use_fused_batchnorm: Enable fused batchnorm.

  Returns:
    An `arg_scope` to use for the inception v3 model.
  """
  batch_norm_params = {
      # Decay for the moving averages.
      'decay': 0.9997,
      # epsilon to prevent 0s in variance.
      'epsilon': 0.001,
      # collection containing update_ops.
      'updates_collections': ops.GraphKeys.UPDATE_OPS,
      # Use fused batch norm if possible.
      'fused': use_fused_batchnorm,
      # collection containing the moving mean and moving variance.
      'variables_collections': {
          'beta': None,
          'gamma': None,
          'moving_mean': [batch_norm_var_collection],
          'moving_variance': [batch_norm_var_collection],
      }
  }

  # Set weight_decay for weights in Conv and FC layers.
  with arg_scope(
      [layers.conv2d, layers_lib.fully_connected],
      weights_regularizer=regularizers.l2_regularizer(weight_decay)):
    with arg_scope(
        [layers.conv2d],
        weights_initializer=init_ops.truncated_normal_initializer(
            stddev=stddev),
        activation_fn=nn_ops.relu,
        normalizer_fn=layers_lib.batch_norm,
        normalizer_params=batch_norm_params) as sc:
      return sc

def _create_partition_checkpoints(sess, checkpoint_dir):
  checkpoint_prefix = os.path.join(checkpoint_dir, "model")
  checkpoint_state_name = "checkpoint"
  v1 = variable_scope.get_variable(
      name="var1",
      shape=[100, 100],
      initializer=init_ops.truncated_normal_initializer(0.5),
      partitioner=partitioned_variables.min_max_variable_partitioner(
          max_partitions=5, axis=0, min_slice_size=8 << 10))
  sess.run(variables.global_variables_initializer())
  v1_value = sess.run(v1._get_variable_list())
  saver = saver_lib.Saver()
  saver.save(
      sess,
      checkpoint_prefix,
      global_step=0,
      latest_filename=checkpoint_state_name)
  return v1_value

def _conv(args, output_size, filter_size, stddev=0.001, bias=True, bias_start=0.0, scope=None):
  if args is None or (nest.is_sequence(args) and not args):
    raise ValueError("`args` must be specified")
  if not nest.is_sequence(args):
    args = [args]

  # Calculate the total size of arguments on dimension 3.
  # (batch_size x height x width x arg_size)
  total_arg_size = 0
  shapes = [a.get_shape().as_list() for a in args]
  height = shapes[0][1]
  width  = shapes[0][2]
  for shape in shapes:
    if len(shape) != 4:
      raise ValueError("Conv is expecting 3D arguments: %s" % str(shapes))
    if not shape[3]:
      raise ValueError("Conv expects shape[3] of arguments: %s" % str(shapes))
    if shape[1] == height and shape[2] == width:
      total_arg_size += shape[3]
    else :
      raise ValueError("Inconsistent height and width size in arguments: %s" % str(shapes))
  
  with vs.variable_scope(scope or "Conv"):
    kernel = vs.get_variable("Kernel", 
      [filter_size[0], filter_size[1], total_arg_size, output_size],
      initializer=init_ops.truncated_normal_initializer(stddev=stddev))
    
    if len(args) == 1:
      res = tf.nn.conv2d(args[0], kernel, [1, 1, 1, 1], padding='SAME')
    else:
      res = tf.nn.conv2d(array_ops.concat(3, args), kernel, [1, 1, 1, 1], padding='SAME')

    if not bias: return res
    bias_term = vs.get_variable( "Bias", [output_size],
      initializer=init_ops.constant_initializer(bias_start))
  return res + bias_term

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.contrib import layers
from tensorflow.contrib.framework.python.ops import arg_scope
from tensorflow.contrib.layers.python.layers import initializers
from tensorflow.contrib.layers.python.layers import layers as layers_lib
from tensorflow.contrib.layers.python.layers import regularizers
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import variable_scope

trunc_normal = lambda stddev: init_ops.truncated_normal_initializer(0.0, stddev)


def inception_v3_base(inputs,
                      final_endpoint='Mixed_7c',
                      min_depth=16,
                      depth_multiplier=1.0,
                      scope=None):
  """Inception model from http://arxiv.org/abs/1512.00567.

  Constructs an Inception v3 network from inputs to the given final endpoint.
  This method can construct the network up to the final inception block
  Mixed_7c.

  Note that the names of the layers in the paper do not correspond to the names
  of the endpoints registered by this function although they build the same

 def testDuplicatedInitializer(self):
   init = init_ops.truncated_normal_initializer(0.0, 1.0)
   self.assertFalse(duplicated_initializer(self, init, 1))

  def testInitFromPartitionVar(self):
    checkpoint_dir = self.get_temp_dir()
    with self.test_session() as session:
      v1 = _create_partition_checkpoints(session, checkpoint_dir)

    # New graph and session.
    with ops.Graph().as_default() as g:
      with self.test_session(graph=g) as session:
        with variable_scope.variable_scope("some_scope"):
          my1 = variable_scope.get_variable(
              name="my1",
              shape=[100, 100],
              initializer=init_ops.zeros_initializer(),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=8 << 10))
          my1_var_list = my1._get_variable_list()
        # Create another variable with different partitions than the variable in
        # the checkpoint.
        with variable_scope.variable_scope("some_other_scope"):
          my2 = variable_scope.get_variable(
              name="var1",
              shape=[100, 100],
              initializer=init_ops.zeros_initializer(),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=16 << 10))
          my2_var_list = my2._get_variable_list()

        checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
            "scope/var1": "some_scope/my1",
            "scope/": "some_other_scope/"})

        session.run(variables.global_variables_initializer())
        my1_values = session.run(my1_var_list)
        self.assertAllEqual(my1_values, v1)
        my2_values = session.run(my2_var_list)
        # Verify we created different number of partitions.
        self.assertNotEquals(len(my2_values), len(v1))
        # Verify the values were correctly initialized inspite of different
        # partitions.
        full_my2_values = np.concatenate(my2_values, axis=0)
        full_v1_values = np.concatenate(v1, axis=0)
        self.assertAllEqual(full_my2_values, full_v1_values)

    # New graph and session.
    with ops.Graph().as_default() as g:
      with self.test_session(graph=g) as session:
        with variable_scope.variable_scope("some_scope"):
          my1 = variable_scope.get_variable(
              name="my1",
              shape=[100, 100],
              initializer=init_ops.truncated_normal_initializer(0.5),
              partitioner=partitioned_variables.min_max_variable_partitioner(
                  max_partitions=5, axis=0, min_slice_size=8 << 10))
          my1_var_list = my1._get_variable_list()

        checkpoint_utils.init_from_checkpoint(checkpoint_dir,
                                              {"scope/var1": my1_var_list,})

        session.run(variables.global_variables_initializer())
        my1_values = session.run(my1_var_list)
        self.assertAllEqual(my1_values, v1)

def embedding_column(categorical_column,
                     dimension,
                     combiner='mean',
                     initializer=None):
  """TPU embedding_column for `tf.feature_column.embedding_column`.

  Note that the interface for TPU embedding_column is different from the non-TPU
  version. The following args available for the non-TPU version are NOT
  supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable.

  Args:
    categorical_column: A categorical_column returned from
        categorical_column_with_identity,  weighted_categorical_column,
        categorical_column_with_vocabulary_list or
        categorical_column_with_vocabulary_file.
    dimension: An integer specifying dimension of the embedding, must be > 0.
    combiner: A string specifying how to reduce if there are multiple entries
      in a single row. For more information, see
      `tf.feature_column.embedding_column`.
    initializer: A variable initializer function to be used in embedding
      variable initialization. If not specified, defaults to
      `tf.truncated_normal_initializer` with mean `0.0` and standard deviation
      `1/sqrt(dimension)`.

  Returns:
    A  _TPUEmbeddingColumn.

  Raises:
    ValueError: if `dimension` not > 0.
    ValueError: if `initializer` is specified but not callable.
  """
  if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS):
    raise TypeError(
        'categorical_column for tpu '
        ' embedding_column must be type %s, got %s.' % (' or '.join([
            cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS
        ]), type(categorical_column)))
  if (dimension is None) or (dimension < 1):
    raise ValueError('Invalid dimension {}.'.format(dimension))

  if (initializer is not None) and (not callable(initializer)):
    raise ValueError('initializer must be callable if specified. '
                     'Embedding of column_name: {}'.format(
                         categorical_column.name))
  if initializer is None:
    initializer = init_ops.truncated_normal_initializer(
        mean=0.0, stddev=1 / math.sqrt(dimension))

  embedding_shape = categorical_column._num_buckets, dimension  # pylint: disable=protected-access

  def _creator(weight_collections, scope):
    embedding_column_layer = fc._EmbeddingColumnLayer(
        embedding_shape=embedding_shape,
        initializer=initializer,
        weight_collections=weight_collections,
        trainable=True,
        name='embedding_column_layer')
    return embedding_column_layer(None, scope=scope)  # pylint: disable=not-callable

  column = _TPUEmbeddingColumn(
      categorical_column=categorical_column,
      dimension=dimension,
      combiner=combiner,
      layer_creator=_creator,
      ckpt_to_load_from=None,
      tensor_name_in_ckpt=None,
      max_norm=None,
      trainable=True)
  # For Embedding column, the initializer is hidden inside the creator Fn, which
  # is not accessiable later. So, we attach it to a speicial field. Also note
  # that non-TPU Embedding column and non-TPU shared Embedding column handle the
  # initializer differently. See shared_embedding_columns for details.
  column._tpu_initializer = initializer
  return column

def _load_embedding_initializer(ckpt_path,
                                embedding_tensor_name,
                                new_vocab_size,
                                embedding_dim,
                                old_vocab_file,
                                new_vocab_file,
                                old_vocab_size=-1,
                                num_oov_buckets=0,
                                initializer=None,
                                max_rows_in_memory=-1):
  """Returns a variable initializer for loading pre-trained embeddings.

  Wrapper around `load_and_remap_matrix_initializer()` specialized for loading
  embedding weights and remapping according to the provided vocab files. See
  docs for `load_and_remap_matrix_initializer()` for more details.

  NOTE: Only for use with div-partitioned variables / vocabularies.

  Args:
    ckpt_path: Path to the TensorFlow checkpoint (version 2, `TensorBundle`)
      from which the old matrix `Tensor` will be loaded.
    embedding_tensor_name: Name of the 2-D `Tensor` to load from checkpoint.
    new_vocab_size: Number of entries in the new vocab.
    embedding_dim: `int` specifying the dimension of the embedding vectors from
      the checkpoint. Must match the number of columns in the old embedding
      matrix.
    old_vocab_file: A scalar `Tensor` of type `string` containing the
      path to the old vocabulary file.
    new_vocab_file: A scalar `Tensor` of type `string` containing the
      path to the new vocabulary file.
    old_vocab_size: The number of entries to consider in the old vocabulary.
      With the default value of -1, the entire old row vocabulary file will be
      used.  Otherwise, only the first `old_vocab_size` entries will be
      considered for remapping.Must be smaller than the length of
      `old_row_vocab_file`.
    num_oov_buckets: `int` specifying the number of out-of-vocabulary
      buckets to use. Must be >= 0.
    initializer: Initializer function that accepts a 1-D tensor as the arg to
      specify the shape of the returned tensor. If `None`, defaults to using
      `truncated_normal_initializer()`.
    max_rows_in_memory: `int` specifying the maximum number of rows to load from
      the checkpoint at once. If less than or equal to 0, the entire matrix will
      be loaded into memory. Setting this arg trades increased disk reads for
      lower memory usage.

  Returns:
    A variable initializer function.
  """
  if initializer is None:
    # TODO(b/25671353): This should be kept in sync with the stddev used by
    # feature_column.py's _EmbeddingColumn.
    initializer = init_ops.truncated_normal_initializer(
        stddev=1.0 / math.sqrt(embedding_dim))

  return _load_and_remap_matrix_initializer(
      ckpt_path=ckpt_path,
      old_tensor_name=embedding_tensor_name,
      new_row_vocab_size=new_vocab_size,
      new_col_vocab_size=embedding_dim,
      old_row_vocab_size=old_vocab_size,
      old_row_vocab_file=old_vocab_file,
      new_row_vocab_file=new_vocab_file,
      old_col_vocab_file=None,
      new_col_vocab_file=None,
      num_row_oov_buckets=num_oov_buckets,
      num_col_oov_buckets=0,
      initializer=initializer,
      max_rows_in_memory=max_rows_in_memory)

def embedding_column(categorical_column,
                     dimension,
                     combiner='mean',
                     initializer=None,
                     max_sequence_length=0,
                     partition_strategy='div'):
  """TPU embedding_column for `tf.feature_column.embedding_column`.

  Note that the interface for TPU embedding_column is different from the non-TPU
  version. The following args available for the non-TPU version are NOT
  supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable.

  Args:
    categorical_column: A categorical_column returned from
        categorical_column_with_identity, weighted_categorical_column,
        categorical_column_with_vocabulary_file,
        categorical_column_with_vocabulary_list,
        sequence_categorical_column_with_identity,
        sequence_categorical_column_with_vocabulary_file,
        sequence_categorical_column_with_vocabulary_list
    dimension: An integer specifying dimension of the embedding, must be > 0.
    combiner: A string specifying how to reduce if there are multiple entries
      in a single row for a non-sequence column. For more information, see
      `tf.feature_column.embedding_column`.
    initializer: A variable initializer function to be used in embedding
      variable initialization. If not specified, defaults to
      `tf.compat.v1.truncated_normal_initializer` with mean `0.0` and
      standard deviation `1/sqrt(dimension)`.
    max_sequence_length: An non-negative integer specifying the max sequence
      length. Any sequence shorter then this will be padded with 0 embeddings
      and any sequence longer will be truncated. This must be positive for
      sequence features and 0 for non-sequence features.
    partition_strategy: Determines how tensors are sharded on the tpu hosts. See
      `tf.nn.safe_embedding_lookup_sparse` for more details. Allowed value are
      `"div"` and `"mod"'. If `"mod"` is used, evaluation and exporting the
      model to CPU will not work. In order to do this, you must shuffle the
      embedding tensors into a single shard.

  Returns:
    A  _TPUEmbeddingColumn.

  Raises:
    ValueError: if `dimension` not > 0.
    ValueError: if `initializer` is specified but not callable.
  """
  if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS):
    raise TypeError(
        'categorical_column for tpu '
        ' embedding_column must be type %s, got %s.' % (' or '.join([
            cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS
        ]), type(categorical_column)))
  if (dimension is None) or (dimension < 1):
    raise ValueError('Invalid dimension {}.'.format(dimension))

  if (initializer is not None) and (not callable(initializer)):
    raise ValueError('initializer must be callable if specified. '
                     'Embedding of column_name: {}'.format(
                         categorical_column.name))
  if initializer is None:
    initializer = init_ops.truncated_normal_initializer(
        mean=0.0, stddev=1 / math.sqrt(dimension))

  embedding_shape = categorical_column._num_buckets, dimension  # pylint: disable=protected-access

  def _creator(weight_collections, scope):
    embedding_column_layer = fc._EmbeddingColumnLayer(
        embedding_shape=embedding_shape,
        initializer=initializer,
        weight_collections=weight_collections,
        trainable=True,
        name='embedding_column_layer')
    return embedding_column_layer(None, scope=scope)  # pylint: disable=not-callable

  column = _TPUEmbeddingColumn(
      categorical_column=categorical_column,
      dimension=dimension,
      combiner=combiner,
      layer_creator=_creator,
      ckpt_to_load_from=None,
      tensor_name_in_ckpt=None,
      max_norm=None,
      trainable=True,
      max_sequence_length=max_sequence_length,
      partition_strategy=partition_strategy)
  # For Embedding column, the initializer is hidden inside the creator Fn, which
  # is not accessiable later. So, we attach it to a speicial field. Also note
  # that non-TPU Embedding column and non-TPU shared Embedding column handle the
  # initializer differently. See shared_embedding_columns for details.
  column._tpu_initializer = initializer
  return column