import torch


def axisangle2quat(axis, angle):
    """
    axis: B x 3: [axis]
    angle: B: [angle]
    returns quaternion: B x 4
    """
    axis = torch.nn.functional.normalize(axis, dim=-1)
    angle = angle.unsqueeze(-1) / 2
    quat = torch.cat([angle.cos(), angle.sin() * axis], dim=-1)
    return quat