Migration Guide#

This guide lists user-visible migrations by release.

Unreleased#

Upgrade PyTorch for compiled COO output#

Applications that request exact COO output from a matrix-backed neighbor method inside torch.compile(fullgraph=True) must upgrade to PyTorch >=2.10. No code change is required for eager execution.

Compiled callers should allocate sufficient matrix capacity instead of relying on NeighborOverflowError: overflow is reported by an asynchronous runtime assertion in a compiled graph. Exact output sizing uses nonzero and may synchronize the host.

JAX Neighbor-List Compilation Boundary#

Use neighbor_list(...) for eager method selection, capacity estimation, allocation, and dispatch. It returns the selected method’s outputs without checking capacity or retrying. Compile a method-specific function such as naive_neighbor_list(...), cell_list(...), or cluster_tile_neighbor_list(...) after choosing the method and capacities. After a compiled call, inspect its matrix counts or fixed-COO recovery metadata before consuming the output. If necessary, enlarge the buffers or recompute stale pair-centric launch metadata, then invoke another specialization.

Replace a compiled unified-dispatch call:

import jax

from nvalchemiops.jax.neighbors import neighbor_list


@jax.jit
def build_neighbors(positions):
    return neighbor_list(
        positions,
        cutoff,
        cell=cell,
        pbc=pbc,
        method="naive",
        max_neighbors=max_neighbors,
    )

with a method-specific compiled call whose allocation and PBC launch metadata are prepared eagerly:

from nvalchemiops.jax.neighbors import (
    compute_naive_num_shifts,
    naive_neighbor_list,
)


shift_range, num_shifts, max_shifts = compute_naive_num_shifts(
    cell, cutoff, pbc
)


@jax.jit
def build_neighbors(positions):
    return naive_neighbor_list(
        positions,
        cutoff,
        cell=cell,
        pbc=pbc,
        max_neighbors=max_neighbors,
        shift_range_per_dimension=shift_range,
        num_shifts_per_system=num_shifts,
        max_shifts_per_system=max_shifts,
    )

Treat cell, pbc, cutoff, allocation sizes, and derived PBC launch metadata as one specialization. Recompute the metadata and create another specialization when cell, pbc, cutoff, or an allocation-driving value changes. Search radii can be JAX arrays; pair-centric cell-list calls additionally close over launch metadata derived from those radii. For compiled fixed COO from naive and cell-list methods, pass a static coo_capacity. The result appends raw required row counts and scalar metadata_valid; compare the counts with pointer differences outside jax.jit, and refresh launch metadata when validity is false. See the Neighbor Lists guide for complete single-system, batched, matrix, fixed-COO, and cluster-tile contracts.

The unreleased aggregate fixed-COO overflow boolean was removed because it could not distinguish row-width shortage, global COO shortage, and stale pair-centric launch metadata. Raw counts preserve the retry size and the clipped pointer preserves exactly what was stored; metadata_valid separately reports whether those counts are trustworthy.

Dual-cutoff calls now reject reversed cutoffs: naive dual-cutoff methods require cutoff2 >= cutoff1, and cluster-tile dual-matrix calls require cutoff2 >= cutoff. Equal cutoffs remain valid.

Cluster-Tile Buffer Capacity#

Cluster-tile construction writes candidate tile pairs to a fixed-size intermediate buffer before producing a neighbor matrix or COO list. Previously, some eager return paths did not check whether construction filled that buffer, while existing guards reported the condition as NeighborOverflowError. All eager Torch and JAX cluster-tile paths now raise TileBufferOverflow when the required tile-pair count exceeds the allocated capacity. The exception provides the required count as num_tiles, the capacity as max_tiles, and the affected system_index for a segmented batch.

The convenience functions continue to estimate an internally allocated buffer when max_tiles_per_group is None. An explicit value that is too small now raises instead of returning incomplete neighbor output. The Tile-buffer capacity section explains how the parameter changes the allocation and how to calculate a retry value from the exception.

JAX transformations require max_tiles_per_group to be a positive static Python integer when the call allocates any tile-index array because it then determines an output shape. Complete caller-supplied tile-index arrays instead determine the capacity; an explicit factor is still validated but never resizes those arrays. Compiled output shapes remain fixed, and a compiled function cannot turn a data-dependent tile count into a Python exception. Adaptive compiled workflows should use the lower-level build function, compare its returned tile counts with the supplied capacities after leaving the compiled region, and run the query only after that check passes.

A segmented eager build reports the first overflowing system. Resizing can therefore require successive retries, and changing segment offsets requires a replacement state initialized with every system marked for rebuild.

Retained Ewald Miller Topology#

Torch and JAX can now retain integer Miller topology separately from Cartesian reciprocal vectors:

miller_indices = generate_ewald_miller_indices(reference_cell, k_cutoff=8.0)
energy = ewald_summation(
    ...,
    cell=current_cell,
    alpha=alpha,
    miller_indices=miller_indices,
)

The generator returns positive-half-space integer rows inside conservative rectangular Miller bounds. It does not filter individual Cartesian-vector magnitudes. The caller owns the topology and must choose bounds that cover all intended cell states. Changing its size changes K and may recompile JAX.

Both Torch and JAX backends also provide ewald_reciprocal_space_from_miller_indices(...). It materializes Cartesian vectors from the supplied cell, then calls the reciprocal component. Torch retains its component-only hybrid_forces option; JAX does not expose it.

JAX ewald_reciprocal_space(...) now follows the tangent carried by k_vectors. Earlier releases dropped that tangent even when vectors were constructed from the differentiated cell, so the cell gradient omitted the reciprocal-vector contribution. This is fixed. A vector is fixed only when it has zero tangent in the active transformation, for example when it was precomputed from a reference cell or passed through jax.lax.stop_gradient(k_vectors). Matching numerical values alone do not create a dependency. Full ewald_summation(k_vectors=...) still treats explicit vectors as fixed metadata.

A reciprocal-component cell derivative holds positions fixed. For a homogeneous-strain derivative, deform positions and cell together.

v0.4.1: Energy Output Layout#

Energy Output Layout (`energy_reduction`)#

Monopole Torch and JAX Ewald, PME, and slab entry points accept keyword-only energy_reduction="atom" | "system" (default "atom").

Mode	Energy shape	Typical use
`"atom"` (default)	`(N,)` per-atom	Weighted atom losses, existing `E.sum()` recipes
`"system"`	`(B,)` per-system	Batched per-system losses without manual `scatter_add`

Direct-output fields are unchanged: forces remain (N, 3), charge gradients remain (N,), and virials remain (B, 3, 3).

For batched per-system losses, prefer energy_reduction="system" over manually reducing per-atom energies with scatter_add or segment_sum:

# Torch: per-system energy and forces from one call.
energy = particle_mesh_ewald(
    positions, charges, cell,
    batch_idx=batch_idx,
    energy_reduction="system",
    neighbor_list=nl, neighbor_ptr=nl_ptr, neighbor_shifts=shifts,
)  # (B,)
forces = -torch.autograd.grad(energy.sum(), positions)[0]

# JAX: same layout under jit when energy_reduction is static.
energy = particle_mesh_ewald(
    positions, charges, cell,
    batch_idx=batch_idx,
    energy_reduction="system",
    neighbor_list=nl, neighbor_ptr=nl_ptr, neighbor_shifts=shifts,
)  # (B,)

On Torch CUDA, eager atom mode may synchronize once per participating component when a materialized uniform cotangent must be proven by value inspection (for example grad_outputs=torch.ones_like(energy) on a contiguous per-atom energy). System mode is structurally sync-free: arbitrary (B,) cotangents reach the cached backward without inspecting atom values. Under torch.compile or CUDA graph capture, only metadata-proven atom cotangents use the fast path; system mode is the guaranteed sync-free layout. JAX adds API/layout parity only; underlying Warp kernels remain atom-buffer-oriented.

v0.4.0: Electrostatics#

Energy-Derivative Training#

For full Ewald/PME APIs, prefer deriving training quantities from the returned energy tensor instead of requesting direct outputs. On the full APIs, each flag below remains functional but emits DeprecationWarning; component APIs such as ewald_real_space, ewald_reciprocal_space, and pme_reciprocal_space keep direct outputs for no-autograd MD/inference loops.

Direct-output flag	Energy-derived replacement
`compute_forces=True`	`forces = -grad(E.sum(), positions)`
`compute_virial=True`	`grad_u = grad(E.sum(), displacement)` with the row-vector displacement recipe; `virial = -grad_u`, `stress = grad_u / V`
`compute_charge_gradients=True`	`dEdq = grad(E.sum(), charges)`
`hybrid_forces=True`	Keep `charges = charge_model(positions)` in the graph and derive forces from energy

Torch full Ewald/PME supports first- and second-order energy derivatives for force/stress training. When a loss mixes forces and stress, take both from a single grad(E.sum(), (positions, strain), create_graph=True) call rather than two separate grad calls – this runs the reciprocal double-backward once instead of twice (see Energy-Derivative Contract). This support is exposed through standard autograd on scalar losses; the electrostatics APIs do not expose public Hessian or Jacobian tensors/functions.

JAX full Ewald/PME supports first-order energy derivatives for positions, charges, and row-vector displacement virials using the same per-system energy-cotangent reducer as Torch. Higher-order JAX support is limited to tested position and charge scalar losses. JAX PME stress/cell/strain, alpha, and precomputed-metadata higher-order paths are unsupported until implemented and tested. JAX direct-output flags remain functional for compatibility in v0.4.0 but are deprecated for differentiable training.

Precomputed Electrostatics Metadata#

Advanced callers can precompute setup-only metadata and pass it to the Ewald/PME entry points instead of regenerating it inside hot loops.

Surface	Precomputed inputs
Ewald reciprocal	`k_vectors`, `miller_bounds`
PME reciprocal	`cell_inv_t`, `volume`, `k_vectors`, `k_squared`, `moduli_x`, `moduli_y`, `moduli_z`
PME B-spline helpers	`compute_bspline_moduli_1d(...)`

These inputs are caches, not differentiable parameters. alpha, cutoffs, mesh controls, batch metadata, neighbor topology, and PME B-spline moduli are treated as constants even if supplied as grad-bearing tensors. For ewald_summation, caller-supplied k_vectors remain static metadata assumed to correspond to the current cell.

The lower-level Torch ewald_reciprocal_space component preserves a supplied k_vectors autograd graph when the cell also requires gradients. Generate those vectors from the same differentiable cell to obtain a physical strain derivative. Detached Cartesian vectors instead define a fixed-k cell derivative; eager execution warns, while the advisory warning is suppressed under torch.compile.

For fixed-cell loops, build metadata once from a detached or stopped-gradient cell and reuse it while the cell is unchanged:

# Torch Ewald fixed-cell loop.
with torch.no_grad():
    k_vectors = generate_k_vectors_ewald_summation(cell, k_cutoff=8.0)
for positions in trajectory:
    energy = ewald_summation(..., cell=cell, k_vectors=k_vectors)

# JAX PME fixed-cell loop.
cell_static = jax.lax.stop_gradient(cell)
cell_inv_t = jnp.linalg.inv(cell_static).transpose(0, 2, 1)
volume = jnp.abs(jnp.linalg.det(cell_static))
reciprocal_cell = 2.0 * jnp.pi * jnp.linalg.inv(cell_static)
k_vectors, k_squared = generate_k_vectors_pme(
    cell_static, mesh_dimensions, reciprocal_cell=reciprocal_cell
)
mesh_nx, mesh_ny, mesh_nz = mesh_dimensions
moduli_x = compute_bspline_moduli_1d(
    jnp.fft.fftfreq(mesh_nx, d=1.0 / mesh_nx), mesh_nx, spline_order
)
moduli_y = compute_bspline_moduli_1d(
    jnp.fft.fftfreq(mesh_ny, d=1.0 / mesh_ny), mesh_ny, spline_order
)
moduli_z = compute_bspline_moduli_1d(
    jnp.fft.rfftfreq(mesh_nz, d=1.0 / mesh_nz), mesh_nz, spline_order
)
for positions in trajectory:
    energy = particle_mesh_ewald(
        positions, charges, cell,
        k_vectors=k_vectors,
        k_squared=k_squared,
        cell_inv_t=cell_inv_t,
        volume=volume,
        moduli_x=moduli_x,
        moduli_y=moduli_y,
        moduli_z=moduli_z,
        mesh_dimensions=mesh_dimensions,
        spline_order=spline_order,
    )

If the cell changes and cell-gradient correctness matters, regenerate the cell-derived metadata for that cell or omit the cache so the wrapper computes it internally. For jax.jit, miller_bounds, mesh_dimensions, spline order, and other shape controls must be concrete static values.

v0.3.0: PyTorch Namespace Migration#

Starting with version 0.3.0, PyTorch is now an optional dependency. The previous PyTorch-based functionality has been moved to a separate nvalchemiops.torch namespace. This section provides a mapping of old import paths to new ones.

Import Path Changes#

Old Import Path	New Import Path
`from nvalchemiops.interactions.dispersion import dftd3`	`from nvalchemiops.torch.interactions.dispersion import dftd3`
`from nvalchemiops.interactions.dispersion import D3Parameters`	`from nvalchemiops.torch.interactions.dispersion import D3Parameters`
`from nvalchemiops.neighbors import neighbor_list`	`from nvalchemiops.torch.neighbors import neighbor_list`
`from nvalchemiops.neighbors import estimate_max_neighbors`	`from nvalchemiops.torch.neighbors.neighbor_utils import estimate_max_neighbors`
`from nvalchemiops.neighborlist import neighbor_list`	`from nvalchemiops.torch.neighbors import neighbor_list`
`from nvalchemiops.neighborlist import cell_list`	`from nvalchemiops.torch.neighbors import cell_list`

Backwards Compatibility#

The old import paths will continue to work but will emit DeprecationWarning messages. They will be removed in a future release.

Naive PBC Metadata Changes#

Advanced callers that precompute periodic metadata for naive neighbor-list methods should update cached arguments as follows:

Old Cached Inputs	New Cached Inputs
`shift_range_per_dimension`, `shift_offset`, `total_shifts`	`shift_range_per_dimension`, `num_shifts_per_system`, `max_shifts_per_system`

The public Torch and JAX APIs now decode periodic shifts on-the-fly inside the neighbor kernels. Materialized shift buffers and shift_offset / total_shifts are no longer part of the public naive-PBC workflow.

Warp Kernels#

If you need direct access to the underlying Warp kernels (without PyTorch), use the non-torch namespaces:

nvalchemiops.neighbors - Warp neighbor list kernels
nvalchemiops.interactions.dispersion - Warp dispersion kernels
nvalchemiops.interactions.electrostatics - Warp electrostatics kernels
nvalchemiops.math - Warp math and spline kernels

These modules comprise both targeted kernels as well as end-to-end launchers where possible, which run the full workflow based on warp.arrays.