import torch


def distortion_derivative(ux: torch.Tensor, uy: torch.Tensor, k: torch.Tensor):
    """Compute the radial distortion derivative term.
    Args:
        ux: The [N x 1] x coordinates of the points in the camera plane.
        uy: The [N x 1] y coordinates of the points in the camera plane.
        k: The radial distortion value.
    Returns:
        The [N x 3 x 3] derivative matrix.
    """
    derivative = torch.zeros((len(ux), 3, 3), dtype=ux.dtype, device=ux.device)
    derivative[:, 0, 0] = 1 + k * (3 * ux ** 2 + uy ** 2)
    derivative[:, 1, 1] = 1 + k * (3 * uy ** 2 + ux ** 2)
    derivative[:, 0, 1] = 2 * k * ux * uy
    derivative[:, 1, 0] = 2 * k * ux * uy
    derivative[:, 2, 2] = 1.0
    return derivative