import torch


def reflect(v: torch.Tensor, axis: torch.Tensor):
    """reflect vector a w.r.t. axis
    Args:
        `v`: tensor of shape `[...,3]`.
        `axis`: tensor with the same shape or dim as `a`.
    Returns:
        the reflected vector
    """
    axis = torch.broadcast_to(axis, v.shape)
    h_vec = 2*axis * torch.sum(axis*v, dim=-1, keepdim=True)
    return h_vec - v