import torch


def grad_fun(x_est,y):
    """
    Defines the gradient function for a denoising problem with White Noise.

    This function demonstrates the use of  Pytorch's autograd to calculate the gradient.
    In this example, the function is equivalent to

    def grad_fun(x_est,y):
        return x_est - y

    :param x_est: Tensor, model estimate of x
    :param y: Tensor, noisy measurements
    :return: grad_x
    """
    # True during training, False during testing
    does_require_grad = x_est.requires_grad

    with torch.enable_grad():
        # Necessary for using autograd
        x_est.requires_grad_(True)
        # Assuming uniform white noise, in the denoising case matrix A is the identity
        error = torch.sum((y - x_est)**2)
        # We retain the graph during training only
        grad_x = torch.autograd.grad(error, inputs=x_est, retain_graph=does_require_grad,
                                     create_graph=does_require_grad)[0]
    # Set requires_grad back to it's original state
    x_est.requires_grad_(does_require_grad)

    return grad_x