import torch


def compute_objective_terms(data, net, need_penalty_terms=False):
    """Construct minibatch of (x,y): compute main objective term and optionally penalty terms.
        Returns expectations over minibatch (last 2 terms only if penalty_terms=True)
        * E[ f(x,y) ]
        * E[ f(x,y)**2 ]
        * [ E[ |df/dx_j|^2 ], E[ |df/dy_j|^2 ] ]
    """
    device = next(net.parameters()).device
    x = data[0].requires_grad_(need_penalty_terms).to(device)
    y = data[1].requires_grad_(False).to(device)
    f = net(x,y)
    E_f = f.mean(0)

    E_f2, E_grad2 = None, None
    if need_penalty_terms:
        E_f2 = (f**2).sum()
        gradx = torch.autograd.grad(f.sum(), x, create_graph=True)[0]
        E_grad2 = (gradx**2).mean(0) # expectation, keep x_j coordinates separate.
    return E_f, E_f2, E_grad2