960f18e759ac0ed76bb9193f29527a209bfd2b1a,qucumber/nn_states/density_matrix.py,DensityMatrix,rotated_gradient,#DensityMatrix#,319

Before Change

            UrhoU = torch.sum(UrhoU_v, dim=(1, 2))

            for i in range(v.shape[0]):
                for j in range(v.shape[0]):
                    rotated_grad[0] += cplx.scalar_mult(
                        UrhoU_v[:, i, j], self.am_grads(v[i, ...], v[j, ...])
                    )
                    rotated_grad[1] += cplx.scalar_mult(
                        UrhoU_v[:, i, j], self.ph_grads(v[i, ...], v[j, ...])
                    )

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

After Change

            ph_grads = self.ph_grads(v, v, expand=True)

            rotated_grad[0] = cplx.einsum("ij,ij...->...", UrhoU_v, am_grads)
            rotated_grad[1] = cplx.einsum("ij,ij...->...", UrhoU_v, ph_grads)

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

Instances