code: def torsion_angles_to_frames(
      r: Rigid,
      alpha: torch.Tensor,
      aatype: torch.Tensor,
      rrgdf: torch.Tensor,
  ):
      # [*, N, 8, 4, 4]
      default_4x4 = rrgdf[aatype, ...]
  
      # [*, N, 8] transformations, i.e.
      #   One [*, N, 8, 3, 3] rotation matrix and
      #   One [*, N, 8, 3]    translation matrix
      default_r = r.from_tensor_4x4(default_4x4)
  
      bb_rot = alpha.new_zeros((*((1,) * len(alpha.shape[:-1])), 2))
      bb_rot[..., 1] = 1
  
      # [*, N, 8, 2]
      alpha = torch.cat(
          [bb_rot.expand(*alpha.shape[:-2], -1, -1), alpha], dim=-2
      )
  
      # [*, N, 8, 3, 3]
      # Produces rotation matrices of the form:
      # [
      #   [1, 0  , 0  ],
      #   [0, a_2,-a_1],
      #   [0, a_1, a_2]
      # ]
      # This follows the original code rather than the supplement, which uses
      # different indices.
  
      all_rots = alpha.new_zeros(
          default_r.get_rots().get_rot_mats().shape
      )
      all_rots[..., 0, 0] = 1
      all_rots[..., 1, 1] = alpha[..., 1]
      all_rots[..., 1, 2] = -alpha[..., 0]
      all_rots[..., 2, 1:] = alpha
  
      all_rots = Rigid(Rotation(rot_mats=all_rots), None)
  
      all_frames = default_r.compose(all_rots)
  
      chi2_frame_to_frame = all_frames[..., 5]
      chi3_frame_to_frame = all_frames[..., 6]
      chi4_frame_to_frame = all_frames[..., 7]
  
      chi1_frame_to_bb = all_frames[..., 4]
      chi2_frame_to_bb = chi1_frame_to_bb.compose(chi2_frame_to_frame)
      chi3_frame_to_bb = chi2_frame_to_bb.compose(chi3_frame_to_frame)
      chi4_frame_to_bb = chi3_frame_to_bb.compose(chi4_frame_to_frame)
  
      all_frames_to_bb = Rigid.cat(
          [
              all_frames[..., :5],
              chi2_frame_to_bb.unsqueeze(-1),
              chi3_frame_to_bb.unsqueeze(-1),
              chi4_frame_to_bb.unsqueeze(-1),
          ],
          dim=-1,
      )
  
      all_frames_to_global = r[..., None].compose(all_frames_to_bb)
  
      return all_frames_to_global
  
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 6 for calling a global function with example usage: torsion_angles_to_frames(r, alpha, aatype, rrgdf) and returns: all_frames_to_global