code: def pseudo_beta_fn(aatype, all_atom_positions, all_atom_masks):
      is_gly = aatype == rc.restype_order["G"]
      ca_idx = rc.atom_order["CA"]
      cb_idx = rc.atom_order["CB"]
      pseudo_beta = torch.where(
          is_gly[..., None].expand(*((-1,) * len(is_gly.shape)), 3),
          all_atom_positions[..., ca_idx, :],
          all_atom_positions[..., cb_idx, :],
      )
  
      if all_atom_masks is not None:
          pseudo_beta_mask = torch.where(
              is_gly,
              all_atom_masks[..., ca_idx],
              all_atom_masks[..., cb_idx],
          )
          return pseudo_beta, pseudo_beta_mask
      else:
          return pseudo_beta
  
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 1 for calling a global function with example usage: pseudo_beta_fn(aatype, all_atom_positions, all_atom_masks) and returns: pseudo_beta,  pseudo_beta_mask, pseudo_beta