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

Before Change

                  of the amplitude and phase RBMS
        :rtype: list[torch.Tensor, torch.Tensor]
        
        rotated_grad = [
            torch.zeros(
                2, getattr(self, net).num_pars, dtype=torch.double, device=self.device
            )
            for net in self.networks
        ]

        Upsi, Upsi_v, v = unitaries.rotate_psi_inner_prod(
            self, basis, sample, include_extras=True
        )

        grad_vp0 = self.rbm_am.effective_energy_gradient(v, reduce=False)
        grad_vp1 = self.rbm_ph.effective_energy_gradient(v, reduce=False)

        // 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
            -cplx.imag(cplx.scalar_divide(rotated_grad[1], Upsi)),  // Imaginary

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