Multi-device execution

Sharding is a build-time value. The mesh spec shards the XC quadrature over grid points and the DF 3-center tensor over auxiliary slabs; the sharded terms are listed for reference; they are constructed by the builder, not by hand.

mesh() composes with either DF backend. The materialized one gives each device a nao²·naux/ndev slab, which is the capacity path until the slab itself stops fitting; past that, df(chunk=...) with mesh= streams instead, each device taking its own slice of the auxiliary range and never holding the tensor at all.

The two shard on different axes. RI-J splits the auxiliary range, since that is what γ is indexed by and all that has to cross devices. RI-K splits the occupied set instead: both E_K and the exchange kernel are sums over occupied orbitals, so each device scans its own slice and the partials are psum-reduced. Padding that axis needs no masking, a zero orbital column contributing nothing to either sum. Both halves of the custom_vjp shard identically, which is what keeps the gradient exact: the exchange Fock the backward returns is the psum of per-device partial kernels, i.e. the single-device matrix. Range-separated exchange rides the same path on the attenuated metric.

The one combination that raises is VV10 (WB97XV) on a sharded grid, and it is a structural limit rather than a gap: its double-grid pair quadrature is nonlocal across shards and cannot be evaluated shard by shard.

dftax.ks.shard.mesh(devices: tuple | list | None = None) -> MeshSpec

Shard the calculation across a 1-D device mesh.

Parameters:

Name	Type	Description	Default
devices	tuple | list | None	the devices to use; None means all local devices (jax.devices()) at build time.	None

Example

KS(mol, xc, coulomb=df("def2-universal-jkfit"), mesh=mesh())
scf_batched(mol, coords_batch, xc, mesh=mesh())   # batch-axis sharding

dftax.ks.shard.MeshSpec

Device-mesh spec (see :func:mesh).

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__dict__ class-attribute

Read-only proxy of a mapping.

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__weakref__ property

list of weak references to the object

__delattr__(name)

__eq__(other)

__hash__()

__init__(devices: tuple | None = None) -> None

__repr__()

__setattr__(name, value)

dftax.ks.terms.ShardedGridXC

XC integral sharded over grid points across a 1-D device mesh.

inner is an ordinary :class:GridXC or :class:StreamedGridXC whose grid-axis arrays are padded to a multiple of the device count and laid out sharded over the mesh (see :func:dftax.ks.shard._pad_shard_grid). The energy runs each device's slice through the unmodified inner math under shard_map (the density is replicated, the partial energies are psum-reduced), so the quadrature is exactly the single-device sum and the collective differentiates natively (autodiff Fock, forces).

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

energy(P)

__init__(inner: XCTerm, devices: tuple) -> None

Initialize self. See help(type(self)) for accurate signature.

dftax.ks.terms.ShardedDFCoulomb

Materialized RI-J with the 3-center tensor sharded over the aux axis.

Each device holds a (nao², naux/ndev) slab of int3c (built directly in shards; see :func:dftax.ks.shard._build_int3c_sharded) and contracts its own slice of γ_P = Σ_μν (μν|P) P_μν; the slices are all_gather-ed (γ is a tiny naux-vector) and the metric quadratic form ½ γᵀ V⁻¹ γ is evaluated replicated. Padded aux columns carry exact zeros in both γ and the (zero-padded) metric inverse, so the energy is the single-device value bit-for-bit up to summation order.

Hybrid exact exchange uses V⁻¹ = LLᵀ and Tr(P K(P)) = Σ_X ⟨P W_X P, W_X⟩ with W = int3c·L: each device builds only its own aux-slab of W (an all-to-all done as ndev rounds of psum, one per destination slab), contracts its slab's exchange partial against the replicated density, and the scalar partials are psum-reduced; per-device memory stays O(nao²·naux/ndev). Padded aux rows are zero in int3c and null in the padded metric, so they contribute exactly nothing to J or K.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

hf_coeff class-attribute

Convert a string or number to a floating-point number, if possible.

hf_coeff_lr class-attribute

Convert a string or number to a floating-point number, if possible.

nao class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

energy(P, S, nocc)

__init__(int3c: Float[Array, 'nao2 nauxp'], int2c_inv: Float[Array, 'nauxp nauxp'], devices: tuple, nao: int = 0, hf_coeff: float = 0.0, int3c_lr: Float[Array, 'nao2 nauxp'] | None = None, int2c_inv_lr: Float[Array, 'nauxp nauxp'] | None = None, hf_coeff_lr: float = 0.0) -> None

Initialize self. See help(type(self)) for accurate signature.

dftax.ks.terms.ShardedStreamedDFCoulomb

Streamed RI-J with the auxiliary axis divided across a device mesh.

The streamed backend never builds the (nao², naux) tensor, which is the only thing that fits at protein scale; this divides its auxiliary axis over the mesh as well, so the per-device work is naux/ndev auxiliary functions streamed chunk at a time. Combining the two is the point: the materialized aux-sharded backend still holds nao²·naux/ndev per device, which for insulin at triple zeta is 15 TiB.

Each device streams its own contiguous slice of the auxiliary index range into its piece of γ, the pieces are all_gather-ed (γ is a naux-vector, so this is negligible traffic), and the metric quadratic form ½ γᵀ V⁻¹ γ is evaluated replicated -- the same decomposition :class:ShardedDFCoulomb uses, and for the same reason: γ is the only quantity that has to cross devices.

The auxiliary basis is replicated, not sliced. The streamed kernel looks an auxiliary function up by index (_eri3c_elem(..., k)) instead of slicing a stored slab, so a device needs the whole (small) auxiliary basis and only its own range of k. Indices are assigned contiguously and the range is padded to a multiple of the device count, with the padding masked to zero, so the gathered γ is already in auxiliary order and truncating it to naux is exact rather than needing a position map.

Hybrids shard too, on a different axis: streamed RI-K sums over occupied orbitals, so each device scans its own slice of them and the partials are psum-reduced (:func:_rik_energy_sharded). Both halves of its custom_vjp shard the same way, so the analytic exchange Fock it returns is unchanged. Range-separated hybrids ride along on the attenuated metric.

__abstractclassvars__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractmethods__ class-attribute

Build an immutable unordered collection of unique elements.

__abstractvars__ class-attribute

Build an immutable unordered collection of unique elements.

__annotations__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ class-attribute

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ class-attribute

__doc__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ class-attribute

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ class-attribute

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ class-attribute

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

hf_coeff class-attribute

Convert a string or number to a floating-point number, if possible.

hf_coeff_lr class-attribute

Convert a string or number to a floating-point number, if possible.

omega class-attribute

Convert a string or number to a floating-point number, if possible.

energy(P, S, nocc)

__init__(basis: BasisData, aux_basis: BasisData, int2c_inv: Float[Array, 'naux naux'], pairs: tuple[Array, Array, Float[Array, npair]] | None, devices: tuple, chunk: int, hf_coeff: float = 0.0, int2c_inv_lr: Float[Array, 'naux naux'] | None = None, hf_coeff_lr: float = 0.0, omega: float = 0.0) -> None

Initialize self. See help(type(self)) for accurate signature.

Across nodes

One process per node runs the same program; distributed() joins them into a process group, after which the mesh spans every GPU of every node and nothing else about the calculation changes. Call it before any other JAX work (dftax itself starts no backend at import, so importing it first is fine).

dftax.ks.distributed.distributed(**kwargs: object) -> tuple[int, int]

Join the JAX process group; returns (process_index, process_count).

Under a multi-task SLURM step (or Open MPI), jax.distributed.initialize discovers the coordinator address, the process count and this process's id from the environment, and each process claims the GPUs its task was given. Outside a multi-task allocation this is a no-op returning (0, 1), so scripts can call it unconditionally.

Parameters:

Name	Type	Description	Default
**kwargs	object	passed straight to jax.distributed.initialize (coordinator_address, num_processes, process_id, local_device_ids); passing any of them forces initialization, which is how a single node can host several processes for testing. local_device_ids is filled in for the one-task-per-node layout, where JAX would otherwise give the task a single GPU.	required

Note

Call this before touching any JAX API that commits the backend (creating an array, querying jax.devices()); JAX cannot join a process group after the local backend is up. Calling it twice is harmless: the second call returns the established group.

dftax.ks.distributed.is_coordinator() -> bool

True in exactly one process of the group: guard printing and file writes with it, never the construction of a term or a solve.

dftax.ks.distributed.barrier(name: str = 'dftax') -> None

Wait until every process in the group reaches this point.

Needed where the processes stop being symmetric: before the coordinator reads a file the others wrote, and at the end of a run, so no process tears down the coordination service while its peers are still talking to it. A no-op in a single-process run.