code: def _axis_angle_rotation(
      axis: str, angle: torch.Tensor
  ) -> torch.Tensor:
      """
      Return the rotation matrices for one of the rotations about an axis
      of which Euler angles describe, for each value of the angle given.
  
      Args:
          axis: Axis label "X" or "Y or "Z".
          angle: any shape tensor of Euler angles in radians
  
      Returns:
          Rotation matrices as tensor of shape (..., 3, 3).
      """
  
      cos = torch.cos(angle)
      sin = torch.sin(angle)
      one = torch.ones_like(angle)
      zero = torch.zeros_like(angle)
  
      if axis == "X":
          R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
      elif axis == "Y":
          R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
      elif axis == "Z":
          R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
      else:
          raise ValueError("letter must be either X, Y or Z.")
  
      return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
  
imports: from typing import Optional
  import torch
  import torch.nn.functional
  from ..common.datatypes import Device
label_desc: Use image node_id 5 for calling a global function with example usage: _axis_angle_rotation(axis, angle) and returns: torch