I have a not totally minimal example, which can be tested in google colab. If you you run it, it shows, that the TPU is set up correctly, but the integration hangs. If you interrupt you get :
tensor([1., 1., 3., 1., 1.], device='xla:1')
xla:1
(0, 1, 2) 0
(0, 2, 1) 1
(1, 0, 2) 1
(1, 2, 0) 0
(2, 0, 1) 0
(2, 1, 0) 1
[[0 1 2 3 4 5]
[0 2 1 3 4 5]
[1 0 2 3 4 5]
[1 2 0 3 4 5]
[2 0 1 3 4 5]
[2 1 0 3 4 5]]
(6, 6)
---------------------------------------------------------------------------
KeyboardInterrupt Traceback (most recent call last)
[<ipython-input-4-2b4f7889adcb>](https://localhost:8080/#) in <cell line: 142>()
140
141
--> 142 plotwf(ppp)
143
144
4 frames
[/usr/local/lib/python3.10/dist-packages/torch/utils/_device.py](https://localhost:8080/#) in __torch_function__(self, func, types, args, kwargs)
60 if func in _device_constructors() and kwargs.get('device') is None:
61 kwargs['device'] = self.device
---> 62 return func(*args, **kwargs)
63
64 # NB: This is directly called from C++ in torch/csrc/Device.cpp
KeyboardInterrupt:
In google colab there are two cells, the first installs TPU support for torch and the needed libs
!pip install cloud-tpu-client==0.10 torch==2.0.0 torchvision==0.15.1 https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-2.0-cp310-cp310-linux_x86_64.whl
!pip install noisyopt
!pip install torchquad
# Python program to compute Hessian in PyTorch
# importing libraries
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.func import hessian
from torchquad import VEGAS, set_up_backend, set_precision
from torch import vmap
import time
from noisyopt import minimizeCompass
from itertools import permutations
from sympy.combinatorics.permutations import Permutation
N_Int_Points = 3000
# set_precision(data_type='float64', backend='torch')
import torch_xla.core.xla_model as xm
dev = xm.xla_device()
torch.set_default_device(dev)
set_up_backend("torch", data_type="float64", torch_enable_cuda=False)
# set_log_level("TRACE")
j = None
# j = torch.complex(torch.tensor(0, dtype=torch.float64), torch.tensor(1, dtype=torch.float64))
ppp = torch.tensor(
[1.0, 1.0, 3.0, 1.0, 1.0]
)
print(ppp)
print(ppp.device)
dist_nuclei = torch.tensor(ppp[0], dtype=torch.float64)
nNuclei = 3
nElectrons = 3
IntElectron = [[-dist_nuclei * (nNuclei//2) * 1.3, dist_nuclei * (nNuclei//2) * 1.3]]
m_Nuclei = 1836
m_Electron = 1
IntNuclei = [[-1.5, 1.5]]
nParticles = nElectrons + nNuclei
m = torch.tensor([m_Nuclei]*nNuclei + [m_Electron]*nElectrons)
offsets = torch.zeros(nParticles) # at 0 there must be high spatial probability density for VEGAS integration to work
for i in range(nNuclei):
offsets[i+nElectrons] = dist_nuclei * (i - nNuclei//2)
def CutRange(x, r):
return torch.sigmoid(x+r)*(1-torch.sigmoid(x-r))
q = torch.tensor([-1]*nElectrons + [1]*nNuclei)
def V(dx):
return torch.exp(-dx**2)
def Vpot(xinp):
"""Potential energy"""
x = xinp + offsets
x1 = x.reshape(-1, 1)
x2 = x.reshape(1, - 1)
dx = x1 - x2
Vdx = q.reshape(-1, 1) * V(dx) * q.reshape(1, -1)
Vdx = Vdx.triu(diagonal=1)
return Vdx.sum()
def Epot(wf, x):
return (torch.conj(wf(x)) * Vpot(x) * wf(x)).real
def H_single(wf, x):
if j is None:
gg = torch.func.grad(lambda x: wf(x).real)(x)
else:
gg = torch.complex(torch.func.grad(lambda x: wf(x).real)(x), torch.func.grad(lambda x: wf(x).imag)(x))
v = 1/(2*m) # from partial integration the minus sign already present
gg = torch.sqrt(v) * gg
return ((torch.dot(torch.conj(gg), gg) + Epot(wf, x)).real)
def H(wf, x):
gg = vmap(lambda x: H_single(wf, x))(x)
return gg
def testwf(ppp, x):
# x = xx[tuple(perms[0]), ] # allows summation over permutations later
return torch.exp(-ppp[1]*x[nElectrons:]**2).prod(-1) * CutRange(x[:nElectrons], ppp[2] * nNuclei//2).prod(-1) * (ppp[4] * torch.sin(x[:nElectrons] * torch.pi / ppp[3])).prod(-1)
def Norm(wf, x):
return (torch.conj(wf(x)) * wf(x)).real
vg = VEGAS()
for i in range(nNuclei):
offsets[i+nElectrons] = ppp[0] * (i - nNuclei//2)
# create permutations
perms = []
perms_p = []
for i in permutations(list(range(nElectrons))):
a = Permutation(list(i))
print(i, a.parity())
perms.append(list(i) + list(range(nElectrons, nParticles)))
perms_p.append(a.parity())
perms = np.array(perms)
print(perms)
print(perms.shape)
def plotwf(ppp):
pl_x = np.linspace(IntElectron[0][0].cpu(), IntElectron[0][1].cpu(), 100)
pl_y = []
pl_y = []
plot_pos = 0
for x in pl_x:
def wf(x):
return testwf(ppp, x)
xinp = [0]*plot_pos + [x] + [0]*(nParticles-1-plot_pos)
xinp = torch.from_numpy(np.array(xinp))
if plot_pos < nElectrons:
int_domain = [IntElectron[0]]*plot_pos + [[x, x+0.01]] + [IntElectron[0]]*(nElectrons-1-plot_pos) + [[-0.1, 0.1]]*nNuclei
else:
int_domain = [IntElectron[0]]*nElectrons + [[-0.1, 0.1]]*(plot_pos-nElectrons) + [[x, x+0.01]] + [[-0.1, 0.1]]*(nParticles - 1 - plot_pos)
integral_value = vg.integrate(lambda y: vmap(lambda y: Norm(lambda x: testwf(ppp, x), y))(y), dim=nParticles, N=10000, integration_domain=int_domain, max_iterations=20)
pl_y.append(integral_value)
pl_y = torch.tensor(pl_y).cpu().numpy()
plt.plot(pl_x + offsets[plot_pos].cpu().numpy(), pl_y)
plt.show()
plotwf(ppp)
def doIntegration(pinp):
start = time.time()
global offsets
ppp = torch.tensor(pinp)
for i in range(nNuclei):
offsets[i+nElectrons] = ppp[0] * (i - nNuclei//2)
IntElectron = [[-ppp[0] * (nNuclei//2) * 1.3, ppp[0] * (nNuclei//2) * 1.3]]
IntNuclei = [[-1.7 / ppp[2], 1.7/ppp[2]]]
Normvalue = integral_value = vg.integrate(lambda y: vmap(lambda y: Norm(lambda x: testwf(ppp, x), y))(y), dim=nParticles, N=N_Int_Points, integration_domain=IntElectron*nElectrons+IntNuclei*nNuclei, max_iterations=30)
integral_value = vg.integrate(lambda y: H(lambda x: testwf(ppp, x), y), dim=nParticles, N=N_Int_Points, integration_domain=IntElectron*nElectrons+IntNuclei*nNuclei, max_iterations=30)
retH = integral_value/Normvalue
print("H", integral_value, ppp, retH, time.time() - start)
return retH.cpu().numpy()
ret = minimizeCompass(doIntegration, x0=ppp.cpu().numpy(), deltainit=0.1, deltatol=0.01, paired=False, bounds=[[0.01, 5]]*ppp.shape[0], errorcontrol=True, funcNinit=10)
print(ret)
Feature
Desired Behavior / Functionality
Setting default device of torch to TPU should work, but it hangs
How Can It Be Tested
I have a not totally minimal example, which can be tested in google colab. If you you run it, it shows, that the TPU is set up correctly, but the integration hangs. If you interrupt you get :
In google colab there are two cells, the first installs TPU support for torch and the needed libs
The second runs the program