import torch


def rot_from_yaw(yaw):
    """
    Get rotation matrix from yaw
    Args:
        yaw: float value representing yaw

    Returns: 3x3 rotation matrix

    """
    if not isinstance(yaw, torch.Tensor):
        yaw = torch.Tensor([yaw])
    cos = torch.cos(yaw)
    sin = torch.sin(yaw)
    z = yaw.new_tensor([0])
    o = yaw.new_tensor([1])

    rot = torch.stack((cos, z, sin, z, o, z, -sin, z, cos)).view(3, 3)
    return rot