netket.optimizer.solver.cholesky_distributed

netket.optimizer.solver.cholesky_distributed(A, b, *, local_tile_size=None, process_grid=None, x0=None)

Solve the linear system using a distributed Cholesky factorization.

This is the distributed/multi-GPU equivalent of cholesky(). It uses jaxmg’s potrs function (LAPACK POTRF+POTRS) which performs a Cholesky factorization followed by triangular solves, optimized for sharded arrays with automatic tiling and communication optimization.

Note

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

Note

JAXMg factorises (the local shard of) A in place. Inside of a jitted function where A is not used after the solve, as in NetKet’s drivers, this needs no copy of A if the default process grid and tile size are used (unless the local shard of A has an awkward size, with no divisor of at least 64 below the tile-size cap). If A is still needed afterwards, for example when calling the solver eagerly or combining it with nan_fallback(), a 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.

Note

This solver expects the input matrix A to be sharded with P(“S”, None).

For single-device or small-scale computations, use cholesky() instead, which has less overhead. Use cholesky_distributed when:

• Running on multiple GPUs with sharded arrays and >= 8k samples.

• NTK/QGT matrix is very large and doesn’t fit on a single device

Parameters:

• A – the matrix A in Ax=b (should be positive definite, sharded)

• b – the vector b in Ax=b

• 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 4096, 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 factorisation 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.

• x0 – unused (kept for API compatibility)

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.cholesky_distributed(local_tile_size=2**12)
>>> driver = nk.driver.VMC_SR(
...     hamiltonian, optimizer, variational_state=vstate,
...     linear_solver=solver, diag_shift=0.01
... )