code: def _embedding_lookup_for_sparse_tensor(
      inp: sparse_tensor.SparseTensor,
      weight: Optional[sparse_tensor.SparseTensor],
      table: tf_variables.Variable,
      feature: tpu_embedding_v2_utils.FeatureConfig,
  ) -> tensor.Tensor:
      """Embedding lookup for sparse tensor based on its feature config.
  
      Args:
        inp: a single SparseTensor input.
        weight: None or SparseTensor which has the same shape of the input.
        table: a table variable.
        feature: a feature config.
  
      Returns:
        Embedding lookup result.
      """
      inp_rank = inp.shape.rank
      # The input rank can be None for sequence input tensor.
      if (
          not feature.output_shape
          and feature.max_sequence_length > 0
          and (inp_rank is None or inp_rank == 2)
      ):
          batch_size = math_ops.cast(
              array_ops.shape(inp)[0], dtype=dtypes.int64
          )
          sparse_shape = array_ops_stack.stack(
              [batch_size, feature.max_sequence_length], axis=0
          )
          # TPU Embedding truncates sequences to max_sequence_length, and if we
          # don't truncate, scatter_nd will error out if the index was out of
          # bounds.
          truncated_inp = sparse_ops.sparse_slice(
              inp, start=[0, 0], size=sparse_shape
          )
  
          dense_output_shape = array_ops_stack.stack(
              [
                  batch_size,
                  feature.max_sequence_length,
                  feature.table.dim,
              ],
              axis=0,
          )
          return array_ops.scatter_nd(
              truncated_inp.indices,
              array_ops.gather(
                  table.read_value(), truncated_inp.values
              ),
              dense_output_shape,
          )
      else:
          if feature.max_sequence_length > 0:
              logging.warning(
                  (
                      "max_sequence_length setting will be ignored because the rank of"
                      " the input tensor is %d which is not 2."
                  ),
                  inp_rank,
              )
          if (
              not feature.validate_weights_and_indices
              and inp_rank is not None
              and inp_rank <= 2
          ):
              return embedding_ops.embedding_lookup_sparse_v2(
                  table,
                  inp,
                  sp_weights=weight,
                  combiner=feature.table.combiner,
              )
          else:
              return embedding_ops.safe_embedding_lookup_sparse_v2(
                  table,
                  inp,
                  sparse_weights=weight,
                  combiner=feature.table.combiner,
              )
  
imports: import functools
  from typing import Any, Dict, Iterable, Optional, Union
  from absl import logging
  from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
  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.framework.constant_op import constant
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  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_embedding_v3_utils
  from tensorflow.python.types import core
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 3 for calling a global function with example usage: _embedding_lookup_for_sparse_tensor(inp, weight, table, feature) and returns: array_ops, embedding_ops, embedding_ops