gpu4pyscf 1.3.0__tar.gz → 1.3.2__tar.gz

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/CHANGELOG

@@ -1,3 +1,22 @@
+v1.3.2 (2025-03-10)
+-------------------
+* Improvements
+  - Dump xc info and grids into to log file
+  - Optimize 4-center integral evalulation CUDA kernels using warp divergent algorithm
+  - Support up to I orbitals in DFT
+  - Fix out-of-bound issue in DFT hessian for heavy atoms (>=19)
+* Deprecation
+  - SM60 is not supported in PyPi package
+
+v1.3.1 (2025-02-04)
+-------------------
+* New Features
+  - Analytical Hessian for PCM solvent model
+  - Driver for 3c methods (wB97x-3c, R2Scan-3c, B97-3c, etc.)
+* Improvements
+  - Preconditioner and computation efficiency of Davidson iterations for TDDFT
+
+
 v1.3.0 (2025-01-07)
 -------------------
 * New Features

{gpu4pyscf-1.3.0/gpu4pyscf.egg-info → gpu4pyscf-1.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 1.3.0
+Version: 1.3.2
 Summary: GPU extensions for PySCF
 Author: PySCF developers
 License: Apache-2.0

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/__config__.py

@@ -14,11 +14,11 @@
 
 import cupy
 
-_num_devices = cupy.cuda.runtime.getDeviceCount()
+num_devices = cupy.cuda.runtime.getDeviceCount()
 
 # TODO: switch to non_blocking stream (currently blocked by libxc)
-_streams = [None] * _num_devices
-for device_id in range(_num_devices):
+_streams = [None] * num_devices
+for device_id in range(num_devices):
     with cupy.cuda.Device(device_id):
         _streams[device_id] = cupy.cuda.stream.Stream(non_blocking=False)
 
@@ -38,11 +38,16 @@ if props['totalGlobalMem'] < 16 * GB:
 mem_fraction = 0.9
 cupy.get_default_memory_pool().set_limit(fraction=mem_fraction)
 
+if props['sharedMemPerBlockOptin'] > 65536:
+    shm_size = props['sharedMemPerBlockOptin']
+else:
+    shm_size = props['sharedMemPerBlock']
+
 # Check P2P data transfer is available
 _p2p_access = True
-if _num_devices > 1:
-    for src in range(_num_devices):
-        for dst in range(_num_devices):
+if num_devices > 1:
+    for src in range(num_devices):
+        for dst in range(num_devices):
             if src != dst:
                 can_access_peer = cupy.cuda.runtime.deviceCanAccessPeer(src, dst)
                 _p2p_access &= can_access_peer

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = '1.3.0'
+__version__ = '1.3.2'
 
 from . import lib, grad, hessian, solvent, scf, dft

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/df.py

@@ -23,13 +23,15 @@ from pyscf.df import df, addons, incore
 from gpu4pyscf.lib.cupy_helper import (cholesky, tag_array, get_avail_mem, 
                                        cart2sph, p2p_transfer, copy_array)
 from gpu4pyscf.df import int3c2e, df_jk
+from gpu4pyscf.df import int3c2e_bdiv
 from gpu4pyscf.lib import logger
 from gpu4pyscf import __config__
-from gpu4pyscf.__config__ import _streams, _num_devices
+from gpu4pyscf.__config__ import _streams, num_devices
 
 MIN_BLK_SIZE = getattr(__config__, 'min_ao_blksize', 128)
 ALIGNED = getattr(__config__, 'ao_aligned', 32)
 GB = 1024*1024*1024
+INT3C2E_V2 = False
 
 LINEAR_DEP_THR = incore.LINEAR_DEP_THR
 GROUP_SIZE = 256
@@ -82,6 +84,42 @@ class DF(lib.StreamObject):
         if auxmol is None:
             self.auxmol = auxmol = addons.make_auxmol(mol, self.auxbasis)
 
+        if INT3C2E_V2:
+            self.intopt = intopt = int3c2e_bdiv.Int3c2eOpt(mol, auxmol)
+            self._cderi = {}
+            self._cderi[0] = _cholesky_eri_bdiv(intopt, omega=omega)
+            ao_pair_mapping = intopt.create_ao_pair_mapping(cart=mol.cart)
+            rows, cols = divmod(cupy.asarray(ao_pair_mapping), mol.nao)
+            intopt.cderi_row = rows
+            intopt.cderi_col = cols
+
+            # intopt.cderi_diag stores the indices for cderi_row that
+            # corresponds to the diagonal blocks. Note this index array can
+            # contain some of the off-diagonal elements which happen to be the
+            # off-diagonal elements while within the diagonal blocks.
+            uniq_l = intopt.uniq_l_ctr[:,0]
+            if mol.cart:
+                nf = (uniq_l + 1) * (uniq_l + 2) // 2
+            else:
+                nf = uniq_l * 2 + 1
+            n_groups = len(uniq_l)
+            ij_tasks = ((i, j) for i in range(n_groups) for j in range(i+1))
+            nbas = intopt.sorted_mol.nbas
+            offset = 0
+            cderi_diag = []
+            for (i, j), bas_ij_idx in zip(ij_tasks, intopt.shl_pair_idx):
+                nfi = nf[i]
+                nfj = nf[j]
+                if i == j: # the diagonal blocks
+                    ish, jsh = divmod(bas_ij_idx, nbas)
+                    idx = np.where(ish == jsh)[0]
+                    addr = offset + idx[:,None] * (nfi*nfi) + np.arange(nfi*nfi)
+                    cderi_diag.append(addr.ravel())
+                offset += bas_ij_idx.size * nfi * nfj
+            intopt.cderi_diag = cupy.asarray(np.hstack(cderi_diag))
+            log.timer_debug1('cholesky_eri', *t0)
+            return self
+
         if omega and omega > 1e-10:
             with auxmol.with_range_coulomb(omega):
                 j2c_cpu = auxmol.intor('int2c2e', hermi=1)
@@ -218,7 +256,7 @@ def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
         # CDERI will be equally distributed to the devices
         # Other devices usually have more memory available than Device 0
         # CDERI will use up to 40% of the available memory
-        use_gpu_memory = naux * npairs * 8 < 0.4 * avail_mem * _num_devices
+        use_gpu_memory = naux * npairs * 8 < 0.4 * avail_mem * num_devices
 
     if use_gpu_memory:
         log.debug("Saving CDERI on GPU")
@@ -226,9 +264,9 @@ def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
         log.debug("Saving CDERI on CPU")
 
     _cderi = {}
-    aux_blksize = (naux + _num_devices - 1) // _num_devices
+    aux_blksize = (naux + num_devices - 1) // num_devices
     aux_blksize = (aux_blksize + ALIGNED - 1) // ALIGNED * ALIGNED
-    for device_id in range(_num_devices):
+    for device_id in range(num_devices):
         p0 = min(aux_blksize*device_id, naux)
         p1 = min(aux_blksize*(device_id+1), naux)
         #for device_id, (p0,p1) in enumerate(lib.prange(0, naux, aux_blksize)):
@@ -246,16 +284,16 @@ def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
     npairs_per_ctr = np.array(npairs_per_ctr)
     total_task_list = np.argsort(npairs_per_ctr)
     task_list_per_device = []
-    for device_id in range(_num_devices):
-        task_list_per_device.append(total_task_list[device_id::_num_devices])
+    for device_id in range(num_devices):
+        task_list_per_device.append(total_task_list[device_id::num_devices])
 
     cd_low_f = cupy.array(cd_low, order='F', copy=False)
     cd_low_f = tag_array(cd_low_f, tag=cd_low.tag)
 
     cupy.cuda.get_current_stream().synchronize()
     futures = []
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             task_list = task_list_per_device[device_id]
             future = executor.submit(_cderi_task, intopt, cd_low_f, task_list, _cderi, aux_blksize,
                                      omega=omega, sr_only=sr_only, device_id=device_id)
@@ -352,7 +390,7 @@ def _cderi_task(intopt, cd_low, task_list, _cderi, aux_blksize,
                 for slice_id, (p0,p1) in enumerate(lib.prange(0, naux, aux_blksize)):
                     tmp = cupy.array(cderi_block[p0:p1], order='C', copy=True)
                     copy_array(tmp, _cderi[slice_id][:p1-p0,ij0:ij1])
-            elif _num_devices > 1:
+            elif num_devices > 1:
                 # Multi-GPU case, copy data to other Devices
                 for dev_id, (p0,p1) in enumerate(lib.prange(0, naux, aux_blksize)):
                     # Making a copy for contiguous data transfer
@@ -364,3 +402,18 @@ def _cderi_task(intopt, cd_low, task_list, _cderi, aux_blksize,
                 _cderi[0][:,ij0:ij1] = cderi_block
             t1 = log.timer_debug1(f'transfer data for {cp_ij_id} / {nq} on Device {device_id}', *t1)    
     return
+
+# Generate CDERI using the new int3c2e_bdiv algorithm
+def _cholesky_eri_bdiv(intopt, omega=None):
+    assert isinstance(intopt, int3c2e_bdiv.Int3c2eOpt)
+    assert omega is None
+    eri3c = intopt.int3c2e_bdiv_kernel()
+    if intopt.mol.cart:
+        eri3c = intopt.orbital_pair_cart2sph(eri3c)
+    auxmol = intopt.auxmol
+    j2c = cupy.asarray(auxmol.intor('int2c2e', hermi=1), order='C')
+    cd_low = cholesky(j2c)
+    aux_coeff = cupy.array(intopt.aux_coeff, copy=True)
+    cd_low = solve_triangular(cd_low, aux_coeff.T, lower=True, overwrite_b=True)
+    cderi = cd_low.dot(eri3c.T)
+    return cderi

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/df_jk.py

@@ -26,7 +26,7 @@ from gpu4pyscf.lib.cupy_helper import contract, transpose_sum, reduce_to_device
 from gpu4pyscf.dft import rks, uks, numint
 from gpu4pyscf.scf import hf, uhf
 from gpu4pyscf.df import df, int3c2e
-from gpu4pyscf.__config__ import _streams, _num_devices
+from gpu4pyscf.__config__ import _streams, num_devices
 
 def _pin_memory(array):
     mem = cupy.cuda.alloc_pinned_memory(array.nbytes)
@@ -453,8 +453,8 @@ def get_jk(dfobj, dms_tag, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-
         mo_coeff = intopt.sort_orbitals(mo_coeff, axis=[1])
 
         futures = []
-        with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-            for device_id in range(_num_devices):
+        with ThreadPoolExecutor(max_workers=num_devices) as executor:
+            for device_id in range(num_devices):
                 future = executor.submit(
                     _jk_task_with_mo,
                     dfobj, dms, mo_coeff, mo_occ,
@@ -474,8 +474,8 @@ def get_jk(dfobj, dms_tag, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-
         mo1s = [intopt.sort_orbitals(mo1, axis=[1]) for mo1 in mo1s]
 
         futures = []
-        with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-            for device_id in range(_num_devices):
+        with ThreadPoolExecutor(max_workers=num_devices) as executor:
+            for device_id in range(num_devices):
                 future = executor.submit(
                     _jk_task_with_mo1,
                     dfobj, dms, mo1s, occ_coeffs,
@@ -486,8 +486,8 @@ def get_jk(dfobj, dms_tag, hermi=0, with_j=True, with_k=True, direct_scf_tol=1e-
     # general K matrix with density matrix
     else:
         futures = []
-        with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-            for device_id in range(_num_devices):
+        with ThreadPoolExecutor(max_workers=num_devices) as executor:
+            for device_id in range(num_devices):
                 future = executor.submit(
                     _jk_task_with_dm, dfobj, dms,
                     hermi=hermi, device_id=device_id,

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/grad/jk.py

@@ -18,7 +18,7 @@ import cupy
 from gpu4pyscf.df.int3c2e import get_int3c2e_ip_jk, VHFOpt, _split_tasks
 from gpu4pyscf.lib.cupy_helper import contract, concatenate, reduce_to_device
 from gpu4pyscf.lib import logger
-from gpu4pyscf.__config__ import _streams, _num_devices
+from gpu4pyscf.__config__ import _streams, num_devices
 
 def _jk_task(with_df, dm, orbo, with_j=True, with_k=True, device_id=0):
     '''  # (L|ij) -> rhoj: (L), rhok: (L|oo)
@@ -61,8 +61,8 @@ def get_rhojk(with_df, dm, orbo, with_j=True, with_k=True):
     '''
     futures = []
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _jk_task, with_df, dm, orbo,
                 with_j=with_j, with_k=with_k, device_id=device_id)
@@ -161,12 +161,12 @@ def get_grad_vjk(with_df, mol, auxmol, rhoj_cart, dm_cart, rhok_cart, orbo_cart,
 
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     futures = []
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _jk_ip_task, intopt, rhoj_cart, dm_cart, rhok_cart, orbo_cart, task_list[device_id],
                 with_j=with_j, with_k=with_k, device_id=device_id, omega=omega)

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/hessian/jk.py

@@ -23,7 +23,7 @@ from gpu4pyscf.scf.int4c2e import libgint
 from gpu4pyscf.hessian.jk import _ao2mo
 from gpu4pyscf.lib import logger
 from gpu4pyscf.lib.cupy_helper import contract, cart2sph, reduce_to_device
-from gpu4pyscf.__config__ import _streams, _num_devices
+from gpu4pyscf.__config__ import _streams, num_devices
 
 NROOT_ON_GPU = 7
 
@@ -171,8 +171,8 @@ def get_jk(dfobj, dms_tag, mo_coeff, mocc, hermi=0,
     mo_coeff = [intopt.sort_orbitals(mo, axis=[0]) for mo in mo_coeff]
 
     futures = []
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _jk_task_with_mo1,
                 dfobj, dms, mo_coeff, mo1s, occ_coeffs,
@@ -415,12 +415,12 @@ def get_int3c2e_hjk(intopt, rhoj, rhok, dm0_tag, with_j=True, with_k=True,
     ncp_ij = len(intopt.log_qs)
     tasks = np.array(list(itertools.product(range(ncp_k), range(ncp_ij))))
     task_list = []
-    for device_id in range(_num_devices):
-        task_list.append(tasks[device_id::_num_devices])
+    for device_id in range(num_devices):
+        task_list.append(tasks[device_id::num_devices])
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_ipip_tasks, intopt, task_list[device_id],
                 rhoj, rhok, dm0_tag, orbo, with_j=with_j, with_k=with_k,

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/hessian/rhf.py

@@ -56,7 +56,7 @@ def _hk_ip1_ip1(rhok1_Pko, dm0, mocc_2):
     hk_ao_ao = cupy.zeros([nao,nao,3,3])
     cupy.get_default_memory_pool().free_all_blocks()
     mem_avail = get_avail_mem()
-    blksize = int((mem_avail*0.4/(nao*nao*3*8)/ALIGNED))*ALIGNED
+    blksize = int(((mem_avail-hk_ao_ao.nbytes)*0.4/(nao*nao*3*8)/ALIGNED))*ALIGNED
     for k0, k1 in lib.prange(0,nnz,blksize):
         #rhok1_Pko_kslice = cupy.asarray(rhok1_Pko[k0:k1])
         rhok1_Pko_kslice = copy_array(rhok1_Pko[k0:k1])
@@ -68,8 +68,9 @@ def _hk_ip1_ip1(rhok1_Pko, dm0, mocc_2):
 
         # (10|0)(0|01) without response of RI basis
         rhok1_Pkl_kslice = contract('piox,ko->pikx', rhok1_Pko_kslice, mocc_2)
+        rhok1_Pko_kslice = None
         hk_ao_ao += contract('pikx,pkiy->ikxy', rhok1_Pkl_kslice, rhok1_Pkl_kslice)
-        rhok1_Pkl_kslice = rhok1_Pko_kslice = None
+        rhok1_Pkl_kslice = None
     return hk_ao_ao
 
 def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None,
@@ -177,14 +178,14 @@ def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None, atmls
         assert blksize > 0
         if hessobj.auxbasis_response:
             hk_ao_aux = cupy.zeros([nao,naux,3,3])
-        for i0, i1 in lib.prange(0,nao,blksize):
-            #wk1_Pko_islice = cupy.asarray(wk1_Pko[:,i0:i1])
-            wk1_Pko_islice = copy_array(wk1_Pko[:,i0:i1])
+            for i0, i1 in lib.prange(0,nao,blksize):
+                #wk1_Pko_islice = cupy.asarray(wk1_Pko[:,i0:i1])
+                wk1_Pko_islice = copy_array(wk1_Pko[:,i0:i1])
+
+                #rhok1_Pko = contract('pq,qiox->piox', int2c_inv, wk1_Pko_islice)
+                rhok1_Pko = solve_j2c(wk1_Pko_islice)
+                wk1_Pko_islice = None
 
-            #rhok1_Pko = contract('pq,qiox->piox', int2c_inv, wk1_Pko_islice)
-            rhok1_Pko = solve_j2c(wk1_Pko_islice)
-            wk1_Pko_islice = None
-            if hessobj.auxbasis_response:
                 # (10|0)(1|00)
                 wk_ip2_Ipo = contract('porx,io->pirx', wk_ip2_P__, mocc_2[i0:i1])
                 hk_ao_aux[i0:i1] += contract('piox,pioy->ipxy', rhok1_Pko, wk_ip2_Ipo)
@@ -200,7 +201,7 @@ def _partial_hess_ejk(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None, atmls
                 wk1_I = contract('yqp,piox->qioxy', int2c_ip1, rhok1_Pko)
                 hk_ao_aux[i0:i1] -= contract('qoi,qioxy->iqxy', rhok0_P_I, wk1_I)
                 wk1_I = rhok0_P_I = None
-        rhok1_Pko = None
+                rhok1_Pko = None
         t1 = log.timer_debug1('contract int3c2e_ip1 with int2c_ip1', *t1)
         
         rho2c_0 = contract('pij,qji->pq', rhok0_P__, rhok0_P__)

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/hessian/rks.py

@@ -46,8 +46,8 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     mocc = mo_coeff[:,mo_occ>0]
     dm0 = numpy.dot(mocc, mocc.T) * 2
 
-    if mf.nlc != '':
-        raise NotImplementedError
+    if mf.do_nlc():
+        raise NotImplementedError("2nd derivative of NLC is not implemented.")
 
     omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
     with_k = mf._numint.libxc.is_hybrid_xc(mf.xc)

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/hessian/uks.py

@@ -48,8 +48,8 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     moccb = mo_coeff[1][:,mo_occ[1]>0]
     dm0a = numpy.dot(mocca, mocca.T)
     dm0b = numpy.dot(moccb, moccb.T)
-    if mf.nlc != '':
-        raise NotImplementedError
+    if mf.do_nlc():
+        raise NotImplementedError("2nd derivative of NLC is not implemented.")
 
     omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
     with_k = mf._numint.libxc.is_hybrid_xc(mf.xc)

{gpu4pyscf-1.3.0 → gpu4pyscf-1.3.2}/gpu4pyscf/df/int3c2e.py

@@ -24,7 +24,7 @@ from gpu4pyscf.lib.cupy_helper import (block_c2s_diag, cart2sph, contract, get_a
                                        reduce_to_device, copy_array, transpose_sum)
 from gpu4pyscf.lib import logger
 from gpu4pyscf.gto.mole import basis_seg_contraction
-from gpu4pyscf.__config__ import _num_devices, _streams
+from gpu4pyscf.__config__ import num_devices, _streams
 
 LMAX_ON_GPU = 8
 FREE_CUPY_CACHE = True
@@ -824,11 +824,11 @@ def get_int3c2e_jk(mol, auxmol, dm0_tag, with_k=True, omega=None):
     futures = []
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_jk_task, intopt, task_list[device_id],
                 dm0_tag, orbo, device_id=device_id, omega=omega)
@@ -935,11 +935,11 @@ def get_int3c2e_ip1_vjk(intopt, rhoj, rhok, dm0_tag, aoslices, with_j=True,
 
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_ip1_vjk_task, intopt, task_list[device_id],
                 rhoj, rhok, dm0_tag, orbo, with_j=with_j, with_k=with_k,
@@ -1033,11 +1033,11 @@ def get_int3c2e_ip2_vjk(intopt, rhoj, rhok, dm0_tag, auxslices,
 
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_ip2_vjk_task, intopt, task_list[device_id],
                 rhoj, rhok, dm0_tag, orbo, with_j=with_j,
@@ -1096,7 +1096,7 @@ def get_int3c2e_ip1_wjk(intopt, dm0_tag, with_k=True, omega=None):
 
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     nao = intopt.mol.nao
     naux = intopt.auxmol.nao
@@ -1107,8 +1107,8 @@ def get_int3c2e_ip1_wjk(intopt, dm0_tag, with_k=True, omega=None):
         wk = np.ndarray([naux,nao,nocc,3], dtype=np.float64, order='C', buffer=mem)
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_ip1_wjk_task, intopt, task_list[device_id],
                 dm0_tag, orbo, wk, with_k=with_k, device_id=device_id, omega=omega)
@@ -1156,11 +1156,11 @@ def get_int3c2e_ip2_wjk(intopt, dm0_tag, with_k=True, omega=None):
 
     aux_ao_loc = np.array(intopt.aux_ao_loc)
     loads = aux_ao_loc[1:] - aux_ao_loc[:-1]
-    task_list = _split_tasks(loads, _num_devices)
+    task_list = _split_tasks(loads, num_devices)
 
     cupy.cuda.get_current_stream().synchronize()
-    with ThreadPoolExecutor(max_workers=_num_devices) as executor:
-        for device_id in range(_num_devices):
+    with ThreadPoolExecutor(max_workers=num_devices) as executor:
+        for device_id in range(num_devices):
             future = executor.submit(
                 _int3c2e_ip2_wjk, intopt, task_list[device_id],
                 dm0_tag, orbo, with_k=with_k, device_id=device_id, omega=omega)