76b2432b1c1ee03c80895a5c7d49cde485323102,qucumber/nn_states/complex_wavefunction.py,ComplexWaveFunction,rotated_gradient,#ComplexWaveFunction#Any#Any#Any#,163

Before Change

        // since grad_vp0/1 are real, can just treat the scalar multiplication
        //  and addition as a matrix multiplication
        torch.matmul(Upsi_v, grad_vp0, out=rotated_grad[0])
        torch.matmul(Upsi_v, grad_vp1, out=rotated_grad[1])

        grad = [
            cplx.real(cplx.scalar_divide(rotated_grad[0], Upsi)),  // Real

After Change

        vr = v.reshape(-1, v.shape[-1])
        raw_grads = [
            self.am_grads(vr).reshape(2, *v.shape[:-1], -1),
            self.ph_grads(vr).reshape(2, *v.shape[:-1], -1),
        ]

        rotated_grad = [cplx.einsum("s...,s...g->...g", Upsi_v, g) for g in raw_grads]

        grad = [
            cplx.real(cplx.scalar_divide(rotated_grad[0], Upsi)),
            cplx.real(cplx.scalar_divide(rotated_grad[1], Upsi)),

Instances