class_docstr: The TPUEmbedding mid level API running on TPU without Embedding accelerator.
  
  NOTE: This mid level API is not intended for large embedding table lookup.
  Embedding tables will be replicated across devices rather than sharding
  across them. To do large embedding table lookup, please use the
  `tpu.experimental.embedding.TPUEmbedding` class. This class is an alternative
  way to do embedding lookups when the TPU doesn't support any version of
  embedding feature. See
  `tpu.experimental.tpu_hardware_feature.embedding_feature` for a detailed
  explanation.
  
  This class has to be created under the `TPUStrategy`, Otherwise a RuntimeError
  will be raised.
  ```python
  strategy = tf.distribute.TPUStrategy(...)
  with strategy.scope():
    embedding = tf.tpu.experimental.embedding.TPUEmbeddingV0(
        feature_config=feature_config,
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
  ```
  When creating a distributed dataset that is to be passed to the lookup
  operation a special input option must be specified:
  
  ```python
  distributed_dataset = (
      strategy.distribute_datasets_from_function(
          dataset_fn=...,
          options=tf.distribute.InputOptions(
              experimental_fetch_to_device=False))
  dataset_iterator = iter(distributed_dataset)
  ```
  
  Below is an example of a training and evaluation step:
  
  ```python
  optimizer = tf.keras.optimizers.SGD(0.1)
  
  @tf.function
  def training_step(dataset_iterator, num_steps):
    def tpu_step(embedding_features):
      with tf.GradientTape() as tape:
        tape.watch(embedding.embedding_table.values())
        activation = embedding(embedding_features)
        model_output = model(activations)
        loss = ...  # some function of labels and model_output
  
      embedding_gradients = tape.gradient(loss,
                                          embedding.embedding_table.values())
      optimizer.apply_gradients(list(zip(gradients,
                                mid_level_api.embedding_tables.values())))
      # Insert your model gradient and optimizer application here
  
    for _ in tf.range(num_steps):
      strategy.run(tpu_step, args=(next(dataset_iterator), ))
  
  @tf.function
  def evalution_step(dataset_iterator, num_steps):
    def tpu_step(embedding_features):
      activations = embedding(embedding_features)
      model_output = model(activations)
      # Insert your evaluation code here.
  
    for _ in tf.range(num_steps):
      strategy.run(tpu_step, args=(next(dataset_iterator), ))
  ```
  
  NOTE: The optimizer used here is a Keras optimizer. In order to make the slot
  variable creation stay consistent between Keras optimizers and
  embedding optimizers, the `slot_variable_creation_fn` argument of the
  embedding optimizers has to be passed with the Keras `add_slot` function. Also
  note that the slot names might be slightly different between them.
  
  ```python
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)
  
  def slot_variable_creation_fn(table, slot_names, slot_initializers):
      slots = {}
      for slot, initializer in zip(slot_names, slot_initializers):
        slots[slot] = optimizer.add_slot(table, slot, initializer)
      return slots
  
  embedding_optimizer = tf.experimental.embedding.Adagrad(
      learning_rate=0.1,
      slot_variable_creation_fn=slot_variable_creation_fn)
  
  # Use the embedding optimizer to create mid level api and keras optimizer to
  # apply gradients.
  ```
class_docstr_tok: ["The","TPUEmbedding","mid","level","API","running","on","TPU","without","Embedding","accelerator",".","NOTE",":","This","mid","level","API","is","not","intended","for","large","embedding","table","lookup",".","Embedding","tables","will","be","replicated","across","devices","rather","than","sharding","across","them",".","To","do","large","embedding","table","lookup",",","please","use","the","`","tpu.experimental.embedding.TPUEmbedding","`","class",".","This","class","is","an","alternative","way","to","do","embedding","lookups","when","the","TPU","does","n't","support","any","version","of","embedding","feature",".","See","`","tpu.experimental.tpu_hardware_feature.embedding_feature","`","for","a","detailed","explanation",".","This","class","has","to","be","created","under","the","`","TPUStrategy","`",",","Otherwise","a","RuntimeError","will","be","raised",".","``","`","python","strategy","=","tf.distribute.TPUStrategy","(","...",")","with","strategy.scope","(",")",":","embedding","=","tf.tpu.experimental.embedding.TPUEmbeddingV0","(","feature_config=feature_config",",","optimizer=tf.tpu.experimental.embedding.SGD","(","0.1",")",")","``","`","When","creating","a","distributed","dataset","that","is","to","be","passed","to","the","lookup","operation","a","special","input","option","must","be","specified",":","``","`","python","distributed_dataset","=","(","strategy.distribute_datasets_from_function","(","dataset_fn=","...",",","options=tf.distribute.InputOptions","(","experimental_fetch_to_device=False",")",")","dataset_iterator","=","iter","(","distributed_dataset",")","``","`","Below","is","an","example","of","a","training","and","evaluation","step",":","``","`","python","optimizer","=","tf.keras.optimizers.SGD","(","0.1",")","@","tf.function","def","training_step","(","dataset_iterator",",","num_steps",")",":","def","tpu_step","(","embedding_features",")",":","with","tf.GradientTape","(",")","as","tape",":","tape.watch","(","embedding.embedding_table.values","(",")",")","activation","=","embedding","(","embedding_features",")","model_output","=","model","(","activations",")","loss","=","...","#","some","function","of","labels","and","model_output","embedding_gradients","=","tape.gradient","(","loss",",","embedding.embedding_table.values","(",")",")","optimizer.apply_gradients","(","list","(","zip","(","gradients",",","mid_level_api.embedding_tables.values","(",")",")",")",")","#","Insert","your","model","gradient","and","optimizer","application","here","for","_","in","tf.range","(","num_steps",")",":","strategy.run","(","tpu_step",",","args=","(","next","(","dataset_iterator",")",",",")",")","@","tf.function","def","evalution_step","(","dataset_iterator",",","num_steps",")",":","def","tpu_step","(","embedding_features",")",":","activations","=","embedding","(","embedding_features",")","model_output","=","model","(","activations",")","#","Insert","your","evaluation","code","here",".","for","_","in","tf.range","(","num_steps",")",":","strategy.run","(","tpu_step",",","args=","(","next","(","dataset_iterator",")",",",")",")","``","`","NOTE",":","The","optimizer","used","here","is","a","Keras","optimizer",".","In","order","to","make","the","slot","variable","creation","stay","consistent","between","Keras","optimizers","and","embedding","optimizers",",","the","`","slot_variable_creation_fn","`","argument","of","the","embedding","optimizers","has","to","be","passed","with","the","Keras","`","add_slot","`","function",".","Also","note","that","the","slot","names","might","be","slightly","different","between","them",".","``","`","python","optimizer","=","tf.keras.optimizers.Adagrad","(","learning_rate=0.1",")","def","slot_variable_creation_fn","(","table",",","slot_names",",","slot_initializers",")",":","slots","=","{","}","for","slot",",","initializer","in","zip","(","slot_names",",","slot_initializers",")",":","slots","[","slot","]","=","optimizer.add_slot","(","table",",","slot",",","initializer",")","return","slots","embedding_optimizer","=","tf.experimental.embedding.Adagrad","(","learning_rate=0.1",",","slot_variable_creation_fn=slot_variable_creation_fn",")","#","Use","the","embedding","optimizer","to","create","mid","level","api","and","keras","optimizer","to","#","apply","gradients",".","``","`"]
docstr: Embedding lookup for ragged tensor based on its feature config.
  
  Args:
    inp: a single rank 2 RaggedTensor input.
    weight: None or RaggedTensor which has the same shape of the input.
    table: a table variable.
    feature: a feature config.
  
  Returns:
    Embedding lookup result.
  
  Raises:
    ValueError: if input ragged tensor is not rank 2 or output shape set in
    the feature config doesn't match with the first dim size of the input.
docstr_tok: ["Embedding","lookup","for","ragged","tensor","based","on","its","feature","config",".","Args",":","inp",":","a","single","rank","2","RaggedTensor","input",".","weight",":","None","or","RaggedTensor","which","has","the","same","shape","of","the","input",".","table",":","a","table","variable",".","feature",":","a","feature","config",".","Returns",":","Embedding","lookup","result",".","Raises",":","ValueError",":","if","input","ragged","tensor","is","not","rank","2","or","output","shape","set","in","the","feature","config","does","n't","match","with","the","first","dim","size","of","the","input","."]
code: def _embedding_lookup_for_ragged_tensor(
      self,
      inp: ragged_tensor.RaggedTensor,
      weight: Optional[ragged_tensor.RaggedTensor],
      table: tf_variables.Variable,
      feature: tpu_embedding_v2_utils.FeatureConfig,
  ) -> tensor.Tensor:
      """Embedding lookup for ragged tensor based on its feature config.
  
      Args:
        inp: a single rank 2 RaggedTensor input.
        weight: None or RaggedTensor which has the same shape of the input.
        table: a table variable.
        feature: a feature config.
  
      Returns:
        Embedding lookup result.
  
      Raises:
        ValueError: if input ragged tensor is not rank 2 or output shape set in
        the feature config doesn't match with the first dim size of the input.
      """
      if inp.shape.rank != 2:
          raise ValueError(
              "Only rank 2 ragged tensor is supported, but got rank {}".format(
                  inp.shape.rank
              )
          )
      batch_size = inp.shape[0]
  
      # This computation needs to placed outside of tpu as the size of the row
      # splits and values can change for different batch which can cause
      # the program to re-compile.
      def ragged_to_dense_outside_compilation(
          inp, weight, batch_size, feature
      ):
          if weight is None:
              weight = ragged_tensor.RaggedTensor.from_row_splits(
                  array_ops.ones_like(inp.values, dtype=dtypes.float32),
                  inp.row_splits,
              )
          if (
              not feature.output_shape
              and feature.max_sequence_length > 0
          ):
              inp = inp.to_tensor(
                  shape=(batch_size, feature.max_sequence_length)
              )
              # Ignore weight if it is a sequence feature.
              weight = array_ops.ones_like(inp, dtype=dtypes.float32)
          elif feature.output_shape:
              # Eagerly run the following op as the result as to be a number in
              # order to use it as part of the output shape.
              with ops.init_scope():
                  output_batch_size = math_ops.reduce_prod(
                      feature.output_shape
                  ).numpy()
              # If the output batch size matches the data batch size, treat it as
              # normal ragged input.
              if output_batch_size == batch_size:
                  inp, weight = inp.to_tensor(), weight.to_tensor()
              # If the data batch size is a factor of the output batch size, the
              # divide result will be the sequence length. Ignore the weights and
              # combiner.
              elif (
                  output_batch_size > batch_size
                  and output_batch_size % batch_size == 0
              ):
                  # Pad or truncate in the sequence dimension
                  seq_length = output_batch_size // batch_size
                  inp = inp.to_tensor(shape=(batch_size, seq_length))
                  # Ignore weight if it is a sequence feature.
                  weight = array_ops.ones_like(
                      inp, dtype=dtypes.float32
                  )
              else:
                  raise ValueError(
                      "Output shape set in the FeatureConfig should be the factor of "
                      "the input data batch size. But instead got output shape {}, "
                      "input data batch size {}".format(
                          feature.output_shape, batch_size
                      )
                  )
          else:
              inp, weight = inp.to_tensor(), weight.to_tensor()
          return inp, weight
  
      inp, weight = tpu_replication.outside_compilation(
          ragged_to_dense_outside_compilation,
          inp=inp,
          weight=weight,
          batch_size=batch_size,
          feature=feature,
      )
  
      embeddings = embedding_ops.embedding_lookup_v2(table, inp)
      weight = array_ops.expand_dims(weight, -1)
      embeddings *= weight
  
      if feature.output_shape:
          with ops.init_scope():
              output_batch_size = math_ops.reduce_prod(
                  feature.output_shape
              ).numpy()
          if output_batch_size == batch_size:
              embeddings = self._apply_combiner_to_embeddings(
                  embeddings, weight, feature.table.combiner
              )
          embeddings = array_ops.reshape(
              embeddings,
              shape=feature.output_shape + [feature.table.dim],
          )
      else:
          if feature.max_sequence_length == 0:
              embeddings = self._apply_combiner_to_embeddings(
                  embeddings, weight, feature.table.combiner
              )
      return embeddings
  
imports: from typing import Any, Dict, Iterable, Optional, Text, Union
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework import tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import embedding_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import sparse_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.tpu import tpu_embedding_base
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.tpu import tpu_replication
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 9 for calling the TPUEmbeddingV0 obj's underlying member method code with example usage: obj._embedding_lookup_for_ragged_tensor(inp, weight, table, feature) and returns: embeddings, inp,  weight