nvalchemi.distributed.ShardedBatch

class nvalchemi.distributed.ShardedBatch(mesh, atom_fields, cell, pbc, n_global, partition_mode='spatial', system_fields=None)

A Batch distributed across a 1-D DeviceMesh.

The chemistry-specific subclass of ShardedCollection: it supplies the atomic-data field->policy map (per-atom fields -> PlainShard; cell / pbc are replicated side metadata) and the Batch-packing logic. The generic scatter / local / gather machinery lives on the base.

Per-atom fields are ShardTensor(Shard(0)) of global shape (n_global, ...) with each rank physically holding n_owned rows. Per-system fields (cell, pbc, and everything else in the batch’s system group, e.g. charge / spin) are replicated.

Obtained via from_batch() (scatter from the source rank) and consumed by DistributedModel via local_batch. full_batch() / to_global_batch() gather back when the user wants a whole-system view.

Parameters:

• mesh (DeviceMesh)

• atom_fields (dict[str, Any])

• cell (torch.Tensor)

• pbc (torch.Tensor)

• n_global (int)

• partition_mode (str)

• system_fields (dict[str, torch.Tensor] | None)

atom_fields()

Return a shallow copy of the atom-field ShardTensor dict.

Return type:

dict[str, Any]

static from_batch(batch, mesh, config, src=0, partition_mode='spatial')

Scatter a full Batch from src rank across mesh.

Parameters:

• batch (Batch | None) – Full-system batch on src; None elsewhere.

• mesh (DeviceMesh) – 1-D device mesh for domain parallelism.

• config (DomainConfig) – Domain-decomposition config. Its mesh / cutoff / grid_dims drive the SpatialPartitioner built internally.

• src (int) – The global rank that holds the full batch (default 0).

• partition_mode (str) –

  How to assign atoms to ranks.

  • "spatial" (default) — SpatialPartitioner, required by halo exchange so a rank’s owned atoms’ neighbors live in adjacent ranks.

  • "contiguous_block" — atoms 0..N/W-1 to rank 0, N/W..2N/W-1 to rank 1, and so on. Avoids degenerate partitions on geometries spatial would choke on (1D chains, perfectly cubic lattices on partition boundaries, clusters in oversized cells).

Return type:

ShardedBatch

Notes

Atoms are scattered honoring the chosen partitioner’s rank assignment verbatim (a per-rank point-to-point scatter), not by an even Shard(0) split of batch.positions. A balanced split would silently override the partitioner whenever the assignment isn’t already balanced (e.g. a cluster not centered in the box), placing atoms on ranks that don’t own their spatial domain so halo exchange can’t reach their real neighbors.

full_batch(dst=0)

Gather all shards into a full Batch on rank dst.

All ranks must call this — the underlying send/recv is collective. Returns None on ranks other than dst.

Parameters:

dst (int)

Return type:

Batch | None

property local_batch: Batch

This rank’s owned atoms as a plain Batch.

Calls .to_local() on each ShardTensor field (no communication, no copy — the returned tensors share storage with the shards). In-place mutations on the returned batch’s tensors propagate back to the ShardTensor automatically; for non-in-place replacements, call update_from_batch() to sync.

local_batch_with_edges(edge_properties=None, node_properties=None)

This rank’s owned atoms as a plain Batch, optionally carrying prepared per-edge and/or per-node routing properties.

The graph-parallel path uses this to hand the wrapper an owned-row batch whose neighbour data the framework prepared: for COO models a "neighbor_list" edge property whose senders are global ids and receivers owned-local; for dense-neighbor_matrix models the per-node "neighbor_matrix" / "num_neighbors" / "neighbor_matrix_shifts" (owned receiver rows, global sender ids into the all-gathered node set). Otherwise identical to local_batch.

Parameters:

• edge_properties (dict[str, torch.Tensor] | None)

• node_properties (dict[str, torch.Tensor] | None)

Return type:

Batch

property n_global: int

Total number of atoms across the mesh.

property n_owned: int

Number of atoms owned by this rank — the local shard size, len(positions.to_local()) (both spatial and contiguous_block modes store Shard(0) per-rank rows).

property num_graphs: int

Number of graphs (systems) — replicated across ranks. Currently inferred as 1 for the single-system domain-decomposition case.

property partition_mode: str

"spatial" / "contiguous_block".

Set at from_batch() time. Both shard per-atom fields Shard(0) (each rank holds n_owned rows); they differ only in how the rank assignment is computed (spatial decomposition vs contiguous blocks).

property rank_assignment: Tensor

rank_assignment[g] is the rank that owns global sharded-atom g.

Atoms are in rank-contiguous order after from_batch()’s scatter-sort, so this is a block tensor with each rank’s block sized by that rank’s n_owned. Built by all-gathering per-rank sizes in a single shot.

Type:

(n_global,) int64 tensor

to_global_batch()

Gather all shards into a full Batch on every rank.

Return type:

Batch

update_from_batch(batch)

Sync non-in-place tensor replacements from batch back into the ShardTensor backing storage.

In-place mutations are already reflected automatically because to_local() returns the backing storage. This method rewraps any per-atom field whose identity has changed on the plain batch.

Parameters:

batch (Batch)

Return type:

None