code: def make_atom14_masks_np(batch):
      batch = tree_map(
          lambda n: torch.tensor(n, device=batch["aatype"].device),
          batch,
          np.ndarray,
      )
      out = make_atom14_masks(batch)
      out = tensor_tree_map(lambda t: np.array(t), out)
      return out
  
imports: import numpy
  import torch
  from .None import residue_constants
  from .tensor_utils import tensor_tree_map, tree_map
label_desc: Use image node_id 2 for calling a global function with example usage: make_atom14_masks_np(batch) and returns: out