code: def chi_angle_atom(atom_index: int) -> np.ndarray:
      """Define chi-angle rigid groups via one-hot representations."""
      chi_angles_index = {}
      one_hots = []
  
      for k, v in chi_angles_atoms.items():
          indices = [atom_types.index(s[atom_index]) for s in v]
          indices.extend([-1] * (4 - len(indices)))
          chi_angles_index[k] = indices
  
      for r in restypes:
          res3 = restype_1to3[r]
          one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
          one_hots.append(one_hot)
  
      one_hots.append(
          np.zeros([4, atom_type_num])
      )  # Add zeros for residue `X`.
      one_hot = np.stack(one_hots, axis=0)
      one_hot = np.transpose(one_hot, [0, 2, 1])
  
      return one_hot
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/esm/openfold_utils/residue_constants.py
imports: import collections
  import copy
  import functools
  from importlib import resources
  from typing import List, Mapping, Tuple
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: chi_angle_atom(atom_index) and returns: one_hot