code: def _partitioned_dynamic_embedding_lookup(
      keys: tf.Tensor,
      de_config: de_config_pb2.DynamicEmbeddingConfig,
      embedding_dim: int,
      sigma_dim: int,
      feature_name: typing.Text,
      service_address: typing.Text = "",
      timeout_ms: int = -1,
  ):
      """Partitions the embeddings of a given `keys` into from given keys.
  
      Args:
        keys: A string `Tensor` of shape [batch] or [batch_size, max_length].
        de_config: Proto DynamicEmbeddingConfig that configs the embedding.
        embedding_dim: a non-negative int denoting the dimension for
          embedding.
        sigma_dim: a positive int denoting the dimension for sigma function.
        feature_name: Feature name for the embedding.
        service_address: The address of a knowledge bank service. If empty, the
          value passed from --kbs_address flag will be used instead.
        timeout_ms: Timeout millseconds for the connection. If negative, never
          timout.
  
      Returns:
        A tuple of `Tensor` of shapes
        - [batch_size, embedding_dim] and [batch_size, sigma_dim], if keys is 2D.
        - [batch_size, max_sequence_length, embedding_dim] and
          [batch_size, max_sequence_length, sigma_dim], if keys is 3D.
      Raises:
        ValueError: if keys' dimension is not 1D or 2D.
      """
      config = de_config_pb2.DynamicEmbeddingConfig()
      config.CopyFrom(de_config)
      config.embedding_dimension = sigma_dim + embedding_dim
      emb = de_ops.dynamic_embedding_lookup(
          keys,
          config,
          feature_name,
          service_address=service_address,
          timeout_ms=timeout_ms,
      )
      if sigma_dim == 0:
          return emb, None
  
      # Decompose the embedding into a sigma part and an embedding part.
      ori_shape = _get_shape_as_list(emb)
      if len(ori_shape) == 2:
          sigma_emb = tf.slice(emb, [0, 0], [-1, sigma_dim])
          embedding = tf.slice(emb, [0, sigma_dim], [-1, -1])
      elif len(ori_shape) == 3:
          sigma_emb = tf.slice(emb, [0, 0, 0], [-1, -1, sigma_dim])
          embedding = tf.slice(emb, [0, 0, sigma_dim], [-1, -1, -1])
      else:
          raise ValueError("Unsupported input shape: %d" % ori_shape)
      return embedding, sigma_emb
  
imports: import collections
  import threading
  import typing
  from research.carls import dynamic_embedding_config_pb2
  from research.carls import dynamic_embedding_ops
  from six.moves import range
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: _partitioned_dynamic_embedding_lookup(keys, de_config, embedding_dim, sigma_dim, feature_name, service_address, timeout_ms) and returns: embedding,  sigma_emb, emb,  None