class_docstr: See: https://arxiv.org/pdf/1711.02508.pdf.
class_docstr_tok: ["See",":","https",":","\/\/arxiv.org\/pdf\/1711.02508.pdf","."]
docstr: See also: https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations.
docstr_tok: ["See","also",":","https",":","\/\/en.wikipedia.org\/wiki\/Rotation_matrix","#","Basic_rotations","."]
code: def axis_and_angle_to_rotation_matrix(
      axis_name: str, angle: torch.Tensor, device, dtype
  ):
      """See also: https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations."""
      axis_name = axis_name.lower()
      assert axis_name in ("x", "y", "z")
      sn = torch.sin(angle)
      cs = torch.cos(angle)
      ones = torch.ones_like(sn)
      zeros = torch.zeros_like(sn)
      if axis_name == "x":
          axis = torch.tensor(
              (1.0, 0.0, 0.0), device=device, dtype=dtype
          ).repeat(angle.size())
          rot_m = torch.stack(
              (ones, zeros, zeros, zeros, cs, -sn, zeros, sn, cs), dim=2
          ).view(-1, 3, 3)
      elif axis_name == "y":
          axis = torch.tensor(
              (0.0, 1.0, 0.0), device=device, dtype=dtype
          ).repeat(angle.size())
          rot_m = torch.stack(
              (cs, zeros, sn, zeros, ones, zeros, -sn, zeros, cs), dim=2
          ).view(-1, 3, 3)
      elif axis_name == "z":
          axis = torch.tensor(
              (0.0, 0.0, 1.0), device=device, dtype=dtype
          ).repeat(angle.size())
          rot_m = torch.stack(
              (cs, -sn, zeros, sn, cs, zeros, zeros, zeros, ones), dim=2
          ).view(-1, 3, 3)
      else:
          raise NotImplementedError(
              f"Not prepared for axis with name {axis_name}"
          )
  
      return rot_m, axis
  
imports: import numpy
  import pytest
  import torch
  import kornia
  from kornia.testing import assert_close
label_desc: Use image node_id 2 for calling the TestAngleOfRotations obj's underlying member method code with example usage: obj.axis_and_angle_to_rotation_matrix(axis_name, angle, device, dtype) and returns: rot_m,  axis