code: def axis_rotation_matrix(axis, theta):
      
      [x, y, z] = axis
      c = np.cos(theta)
      s = np.sin(theta)
      return np.array([
          [x**2 + (1 - x**2) * c, x * y * (1 - c) - z * s, x * z * (1 - c) + y * s],
          [x * y * (1 - c) + z * s, y ** 2 + (1 - y**2) * c, y * z * (1 - c) - x * s],
          [x * z * (1 - c) - y * s, y * z * (1 - c) + x * s, z**2 + (1 - z**2) * c]
      ])
docs: Returns a rotation matrix around the given axis