K = 64
axis = torch.linspace(-6, 6, K)
for d in (1, 2, 3):
energy_d = make_energy(d)
pts = torch.stack(torch.meshgrid(*([axis] * d), indexing='ij'),
-1).reshape(-1, d)
with torch.no_grad():
log_Z = (torch.logsumexp(-energy_d(pts), 0)
+ d * math.log(axis[1] - axis[0]))
print(f'd = {d}: {len(pts):>7d} energy evaluations, log Z = {log_Z:7.3f}')
for d in (10, 784):
print(f'd = {d}: {K}^{d} = 10^{d * math.log10(K):.0f} energy evaluations')