code: def frames_and_literature_positions_to_atom14_pos(
      r: Rigid,
      aatype: torch.Tensor,
      default_frames,
      group_idx,
      atom_mask,
      lit_positions,
  ):
      # [*, N, 14]
      group_mask = group_idx[aatype, ...]
  
      # [*, N, 14, 8]
      group_mask = nn.functional.one_hot(
          group_mask,
          num_classes=default_frames.shape[-3],
      )
  
      # [*, N, 14, 8]
      t_atoms_to_global = r[..., None, :] * group_mask
  
      # [*, N, 14]
      t_atoms_to_global = t_atoms_to_global.map_tensor_fn(
          lambda x: torch.sum(x, dim=-1)
      )
  
      # [*, N, 14, 1]
      atom_mask = atom_mask[aatype, ...].unsqueeze(-1)
  
      # [*, N, 14, 3]
      lit_positions = lit_positions[aatype, ...]
      pred_positions = t_atoms_to_global.apply(lit_positions)
      pred_positions = pred_positions * atom_mask
  
      return pred_positions
  
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 7 for calling a global function with example usage: frames_and_literature_positions_to_atom14_pos(r, aatype, default_frames, group_idx, atom_mask, lit_positions) and returns: pred_positions