code: def gather_token_embedding(
      tensor_names, state_dict, tn_pattern, scale=True, emb_dim=None
  ):
      target_tn = []
      for tn in tensor_names:
          if tn_pattern in tn.split("."):
              target_tn.append(tn)
              continue
      target_tensor = [state_dict[name] for name in target_tn]
      target_tensor = np.concatenate(target_tensor, axis=0)
      if emb_dim is None:
          emb_dim = target_tensor.shape[1]
      if scale:
          target_tensor = target_tensor * (emb_dim**0.5)
      return target_tensor, target_tn
  
file_path: lightseq/lightseq/training/ops/pytorch/export.py
imports: from collections import OrderedDict
  import numpy
  from lightseq.training.ops.pytorch.util import get_pos_embedding
  from lightseq.training import LSTransformerEncoderLayer, LSTransformerDecoderLayer
label_desc: Use image node_id 1 for calling a global function with example usage: gather_token_embedding(tensor_names, state_dict, tn_pattern, scale, emb_dim) and returns: target_tensor,  target_tn