Precision and placement

Precision is the dtype policy of a solve, and Layout places experiments and factorizations on the devices of a mesh. See precision and multi-GPU parallelism.

diff_soup.Precision

The dtype policy of a solve, declared once on the problem.

Pass it as precision= when building the problem, for example spec.build(precision=...). Precision() is float64 compute with the "fp64" solve mode, the library's default. "fp32_ir" refines float32 factors against a float64 residual, and "adaptive" switches between it and "fp64" at runtime.

Parameters:

Name	Type	Description	Default
compute_dtype		Floating dtype of states, parameters, residuals, and Jacobians inside the solve.	required
solve_mode		Linear solve mode, one of "fp64", "fp32", "fp32_ir", or "adaptive", with "fp64" read as "fp32" under a float32 compute_dtype.	required
adjoint_compute_dtype		Dtype of the reverse sweep, at least as wide as compute_dtype, which it inherits when None.	required
adjoint_ir_tol		Positive refinement target of the reverse implicit-function solves, inheriting the forward target when None.	required
forward_sensitivity_solve_mode		Solve mode of forward sensitivities, inheriting solve_mode when None.	required
reverse_sensitivity_solve_mode		Solve mode of the reverse sweep, inheriting solve_mode when None.	required
jacobian_values_dtype		Dtype of the Jacobian entries alone, no wider than compute_dtype.	required

fp64 classmethod

fp64() -> 'Precision'

Return the policy of float64 compute with float64 factors, the default.

Returns:

Type	Description
'Precision'	A Precision with float64 compute and the "fp64" solve mode.

fp32 classmethod

fp32() -> 'Precision'

Return the policy of float32 compute with float32 factors, end to end.

States, parameters, factors, and both derivative directions run in float32.

Returns:

Type	Description
'Precision'	A Precision with float32 compute and the "fp32" solve mode.

fp32_factor classmethod

fp32_factor() -> 'Precision'

Return the policy of float64 compute with unrefined float32 factors.

Forward sensitivities also use float32 factors, and the reverse sweep solves in float64.

Returns:

Type	Description
'Precision'	A Precision with float64 compute and the "fp32" solve mode.

fp32_ir classmethod

fp32_ir() -> 'Precision'

Return the policy of float64 compute with float32 factors refined against a float64 residual.

Forward sensitivities refine the same way, and the reverse sweep solves in float64.

Returns:

Type	Description
'Precision'	A Precision with float64 compute and the "fp32_ir" solve mode.

diff_soup.Layout

Placement of experiments and factorizations on the devices of a mesh.

As layout= on ds.lower it splits experiments over the data axis, and on spec.build it factors a problem over the model axis. Settings on a ds.linear.CuDSS object, such as reorder="metis", win over the layout's, and the object's None leaves them to the layout.

Parameters:

Name	Type	Description	Default
mesh		Mesh from global_mesh with axes named "data" or "model", or None for one device.	required
dof_sharding		Stored on the layout and not read by the library.	required
reorder		Fill-reducing ordering of cuDSS, where "metis" selects METIS and the default "amd" keeps cuDSS's own.	required
hybrid_memory		If True, cuDSS keeps part of the factor in host memory, and a model axis across processes takes this from the layout alone.	required
hybrid_device_memory_limit_mb		Positive device memory cap in MB for the hybrid mode, which needs hybrid_memory=True.	required
balance_steps		If True, move experiments between devices by their step counts when a device holds several.	required
processes		Number of processes the run must have, or None for any, and None or 1 when the data axis spans several devices.	required

diff_soup.global_mesh

global_mesh(axis_names: Sequence[str], *, shape: Sequence[int] | None = None) -> Mesh

Build a mesh over every device JAX sees, across all processes.

A one-axis mesh spans every device, and a two-axis mesh needs shape.

Parameters:

Name	Type	Description	Default
axis_names	Sequence[str]	Distinct names from ("data", "model").	required
shape	Sequence[int] | None	Devices per axis, whose product must equal the device count.	None

Returns:

Type	Description
Mesh	A jax.sharding.Mesh for Layout(mesh=...).

diff_soup.mesh_groups

mesh_groups(mesh: Mesh, sizes: Sequence[int]) -> tuple[Mesh, ...]

Split a mesh along its data axis into consecutive submeshes.

Pass the groups as device_groups= to ds.inference.checkpointed_map.

Parameters:

Name	Type	Description	Default
mesh	Mesh	Mesh to split, usually from global_mesh.	required
sizes	Sequence[int]	Number of data devices in each group, positive and summing to the axis width.	required

Returns:

Type	Description
tuple[Mesh, ...]	One mesh per entry of sizes, or the mesh itself when it has no data axis.

---

Multi-process execution

ds.distributed.run splits a study's experiments over processes, and ds.distributed.initialize starts the runtime for a model axis across them.

diff_soup.distributed.run

run(study: Any, **options: Any) -> Iterator[Any]

Lower a study over the processes of a run, and close it when the block ends.

Use it as with ds.distributed.run(study) as lowered:. In a run of several processes, process 0 runs the block and the others evaluate their experiments until it ends, so write results inside the block.

Parameters:

Name	Type	Description	Default
study	Any	The Study to lower.	required
**options	Any	Passed to ds.lower, for example objective= or layout=.	{}

Yields:

Type	Description
Any	The lowered study.

Raises:

Type	Description
TypeError	If study is not a Study.
SystemExit	With status 0, on every process above 0 once process 0 leaves the block.

diff_soup.distributed.initialize

initialize(*, coordinator_address: str | None = None, num_processes: int | None = None, process_id: int | None = None, debug: bool = False, **kwargs: Any) -> None

Start JAX's multi-process runtime for a run of several processes.

Call it before any array is created. Without arguments it detects a Slurm launch with several tasks, python -m diff_soup.distributed.launch, or an MPI launch with DIFFSOUP_COORDINATOR set, and otherwise does nothing.

Parameters:

Name	Type	Description	Default
coordinator_address	str | None	host:port of process 0, given together with num_processes and process_id.	None
num_processes	int | None	Number of processes in the run.	None
process_id	int | None	Rank of this process.	None
debug	bool	If True, check that optimizer decisions agree across processes.	False
**kwargs	Any	Passed to jax.distributed.initialize.	{}

Raises:

Type	Description
ValueError	If only some of coordinator_address, num_processes, and process_id are given.

diff_soup.distributed.configure

configure() -> None

Add diff-soup's default XLA GPU flags to XLA_FLAGS, keeping any flag already set.

It runs when diff_soup is imported and again in initialize.

diff_soup.distributed.is_primary

is_primary() -> bool

Return whether this process is process 0.

Returns:

Type	Description
bool	True on process 0.

diff_soup.distributed.owns

owns(mesh: Mesh) -> bool

Return whether any device of mesh belongs to this process.

Parameters:

Name	Type	Description	Default
mesh	Mesh	Mesh to test.	required

Returns:

Type	Description
bool	True when this process holds a device of mesh.

diff_soup.distributed.local_value

local_value(value: Any, dtype=None) -> np.ndarray

Read a value as a NumPy array on this process.

A global array that this process cannot fully address must be replicated, and its first local shard is read.

Parameters:

Name	Type	Description	Default
value	Any	Array or array-like to read.	required
dtype		Optional NumPy dtype of the result.	None

Returns:

Type	Description
ndarray	A host numpy.ndarray.

Raises:

Type	Description
ValueError	If value is a partitioned global array.

diff_soup.distributed.from_local_data

from_local_data(data: Any, mesh: Mesh, spec: PartitionSpec, *, global_shape=None) -> jax.Array

Assemble a global array from the piece of it each process holds.

Parameters:

Name	Type	Description	Default
data	Any	This process's local piece.	required
mesh	Mesh	Mesh the array is placed on.	required
spec	PartitionSpec	PartitionSpec of the array over mesh.	required
global_shape		Shape of the global array, inferred from the local pieces when None.	None

Returns:

Type	Description
Array	A jax.Array sharded by NamedSharding(mesh, spec).

diff_soup.distributed.key

key(seed: int) -> jax.Array

Return jax.random.key(seed).

Pass the same seed on every process so they draw the same random numbers.

Parameters:

Name	Type	Description	Default
seed	int	Integer seed.	required

Returns:

Type	Description
Array	A typed JAX random key.

diff_soup.distributed.any_process

any_process(flag: Any) -> bool

Return whether flag is true on any process.

It is a collective, so every process taking part must call it. In one process it returns flag.

Parameters:

Name	Type	Description	Default
flag	Any	Boolean, or a scalar array read with local_value.	required

Returns:

Type	Description
bool	True when flag is true on at least one process.