code: def atom14_to_atom37(atom14, batch):
      atom37_data = batched_gather(
          atom14,
          batch["residx_atom37_to_atom14"],
          dim=-2,
          no_batch_dims=len(atom14.shape[:-2]),
      )
  
      atom37_data = atom37_data * batch["atom37_atom_exists"][..., None]
  
      return atom37_data
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/esm/openfold_utils/feats.py
imports: import torch
  import torch.nn
  from .None import residue_constants
  from .rigid_utils import Rigid, Rotation
  from .tensor_utils import batched_gather
label_desc: Use image node_id 2 for calling a global function with example usage: atom14_to_atom37(atom14, batch) and returns: atom37_data