def METHOD_NAME(
    F,
    matrix: Tensor,
    num_data_points: Optional[int] = None,
    float_type=np.float32,
) -> Tensor:
    """
    This function extracts the diagonal of a batch matrix.
    Parameters
    ----------
    F
        A module that can either refer to the Symbol API or the NDArray
        API in MXNet.
    matrix
        matrix of shape (batch_size, num_data_points, num_data_points).
    num_data_points
        Number of rows in the kernel_matrix.
    Returns
    -------
    Tensor
        Diagonals of kernel_matrix of shape (batch_size, num_data_points, 1).
    """
    return F.linalg.gemm2(
        F.broadcast_mul(F.eye(num_data_points, dtype=float_type), matrix),
        F.ones_like(F.slice_axis(matrix, axis=2, begin=0, end=1)),
    )