netket.optimizer.solver.pinv_smooth_distributed

netket.optimizer.solver.pinv_smooth_distributed(A, b, *, rtol=1e-14, rtol_smooth=1e-14, local_tile_size=None, process_grid=None, x0=None)

Solve the linear system using a distributed pseudo-inverse from eigendecomposition.

This is the distributed/multi-GPU equivalent of pinv_smooth(). It uses jaxmg’s syevd function (cuSOLVERMp SYEVD) to compute eigenvalues and eigenvectors of a symmetric/Hermitian matrix in a distributed manner, then applies the same smoothed regularization as pinv_smooth().

The eigenvalues \(\lambda_i\) smaller than \(r_\textrm{tol} \lambda_\textrm{max}\) are truncated (where \(\lambda_\textrm{max}\) is the largest eigenvalue).

The eigenvalues are further smoothed with another filter, originally introduced in Medvidovic, Sels arXiv:2212.11289 (2022), given by the following equation

\[\tilde\lambda_i^{-1}=\frac{\lambda_i^{-1}}{1+\big(\epsilon\frac{\lambda_\textrm{max}}{\lambda_i}\big)^6}\]

Note

Requires jaxmg >= 1.4.1 and one python process per GPU (e.g. launch with djaxrun).

Note

JAXMg overwrites (the local shard of) A with its work space, and stores the eigenvectors in a new buffer as large as A, so the solve needs about twice the memory of A. If A is still needed afterwards, for example when calling the solver eagerly or combining it with nan_fallback(), a further copy of the local shard of A is made first. A non-default process_grid, or a local_tile_size not dividing the local shard, also require a (padded) copy.

Parameters:

• A – the matrix A in Ax=b (should be symmetric/Hermitian, sharded)

• b – the vector b in Ax=b

• rtol (float) – Relative tolerance for small eigenvalues of A. For the purposes of rank determination, eigenvalues are treated as zero if they are smaller than rtol times the largest eigenvalue of A.

• rtol_smooth (float) – Regularization parameter used with a similar effect to rtol but with a softer curve. See \(\epsilon\) in the formula above.

• local_tile_size –

  Tile size for matrix A. Controls the tiling strategy for distributed computation (see choosing a tile size). Defaults to the largest tile dividing the local shard of A (so that A needs no padding), capped at 512, if one of at least 64 exists.

  The tile size determines the memory/communication tradeoff:

  • Larger tiles (closer to matrix size): Less communication overhead, more memory per device

  • Smaller tiles (64-1024): More communication, less memory per device, better load balancing

  See ArXiV:2601.14466 for details specific to NQS. For most applications the default works well. Adjust if you encounter memory issues (decrease local_tile_size) or excessive communication overhead (increase local_tile_size).

• process_grid – Shape (process_rows, process_cols) of the 2D cuSOLVERMp process grid, which must have one slot per device. Defaults to (n_devices, 1), which matches the row sharding NetKet uses for the NTK (use_ntk=True), and therefore requires no redistribution of A (by default, an A sharded by columns may use the (1, n_devices) grid instead). A genuinely 2D grid reduces the communication volume of the eigendecomposition itself, at the cost of redistributing A first, and can be faster for large matrices. It requires the devices of the mesh to be in the order of jax.devices(), as in the mesh set by NetKet.

Returns:

(solution, None) where solution is the unraveled result.

Return type:

tuple

Example

>>> import netket as nk
>>> import jax
>>> from jax.sharding import PartitionSpec as P
>>>
>>> # For multi-GPU setup with sharding enabled
>>> solver = nk.optimizer.solver.pinv_smooth_distributed(
...     rtol=1e-14, rtol_smooth=1e-14, local_tile_size=2**12
... )
>>> driver = nk.driver.VMC_SR(
...     hamiltonian, optimizer, variational_state=vstate,
...     linear_solver=solver, diag_shift=0.01
... )