fix(neighbors): harden selective cluster-tile state

CHANGELOG.md

@@ -1,5 +1,29 @@
 # Changelog
 
+## Unreleased
+
+### Added
+
+- PyTorch cluster-tile selective calls can append caller-owned tile state with
+  `return_state=True`, without changing the default neighbor-list return arity.
+
+### Fixed
+
+- Unbatched JAX naive dual-cutoff PBC neighbor lists now populate both cutoff
+  outputs when using the default `wrap_positions=True`. Previously this path
+  wrapped positions but skipped the fill kernel, leaving zero counts and padded
+  matrices.
+- Batched PyTorch cluster-tile segmented COO validates fixed topology, offsets,
+  counts, and tile-state capacities before launching Warp kernels.
+- Segmented cluster-tile COO kernels now bound writes by physical output capacity
+  when compiled steady-state calls reuse malformed caller metadata, preventing
+  invalid segments from escaping their backing buffers.
+- Compiled unified PyTorch cluster-tile dispatch now rejects tensor-valued PBC
+  rather than treating it as fully periodic. Eagerly validate PBC and compile the
+  direct single-system fixed-state route instead.
+- JAX cluster-tile empty selective rebuilds now preserve false-flag state and
+  clear true-flag pair and tile counts while retaining fixed-capacity storage.
+
 ## 0.4.0 - 2026-07-13
 
 ### Added

docs/userguide/components/neighborlist.md

@@ -1245,6 +1245,133 @@ segmented-COO outputs in both the PyTorch and JAX `batch_naive` / `batch_cell_li
 paths (single-system paths take a whole-system flag of shape `(1,)`). It is not
 combined with differentiable per-pair geometry.
 
+Cluster-tile selective rebuilds also require the tile-list buffers. JAX returns that
+state as part of every selective result. Zero-atom calls keep the same arity as
+nonempty calls: selective matrix output contains the matrix triple plus tile state
+(and both matrix triples for dual cutoff), while segmented COO remains the normal
+seven-array single-system or ten-array batched tuple.
+
+PyTorch keeps the existing return arity by default. To receive the tile state, select
+the cluster-tile method explicitly and pass `return_state=True` with
+`rebuild_flags`:
+
+- Single-system matrix output appends `(num_tiles, tile_row_group,
+  tile_col_group)`, yielding six tensors for one cutoff and nine for two.
+- Single-system segmented COO returns `(neighbor_list, pair_offsets,
+  pair_counts, neighbor_list_shifts)` and appends `(num_tiles,
+  tile_row_group, tile_col_group)` with `return_state=True`, yielding seven
+  tensors. The fixed buffers have shapes `(2, max_pairs)`, `(2,)`, `(1,)`, and
+  `(max_pairs, 3)`; `pair_offsets` is `[0, max_pairs]`, and
+  `pair_counts[0]` gives the active prefix. Values after that prefix are
+  inactive.
+- Batched matrix output appends `(tile_offsets, tile_counts, num_tiles,
+  tile_row_group, tile_col_group, tile_system)`, yielding nine tensors for one
+  cutoff and twelve for two.
+- Batched segmented COO appends the same six state tensors to its four topology
+  outputs, yielding ten tensors.
+
+```python
+from nvalchemiops.torch.neighbors import neighbor_list
+
+(
+    neighbor_matrix,
+    num_neighbors,
+    shifts,
+    tile_offsets,
+    tile_counts,
+    num_tiles,
+    tile_row_group,
+    tile_col_group,
+    tile_system,
+) = neighbor_list(
+    positions,
+    cutoff,
+    cell=cells,
+    pbc=pbc,
+    batch_ptr=batch_ptr,
+    method="batch_cluster_tile",
+    rebuild_flags=rebuild_flags,
+    return_state=True,
+)
+```
+
+For a single-system fixed-capacity COO workflow, retain the public topology
+buffers and tile state from the previous step:
+
+```python
+(
+    neighbor_list_coo,
+    pair_offsets,
+    pair_counts,
+    shifts_coo,
+    num_tiles,
+    tile_row_group,
+    tile_col_group,
+) = neighbor_list(
+    positions,
+    cutoff,
+    cell=cell,
+    pbc=pbc,
+    method="cluster_tile",
+    return_neighbor_list=True,
+    rebuild_flags=rebuild_flags,  # shape (1,)
+    neighbor_list=neighbor_list_coo,
+    pair_offsets=pair_offsets,  # tensor([0, max_pairs], dtype=torch.int32)
+    pair_counts=pair_counts,
+    neighbor_list_shifts=shifts_coo,
+    num_tiles=num_tiles,
+    tile_row_group=tile_row_group,
+    tile_col_group=tile_col_group,
+    return_state=True,
+)
+```
+
+Pass the returned topology buffers and state back on the next call. `return_state`
+is intentionally rejected for auto-selected and non-cluster-tile methods, and it
+requires `rebuild_flags`. Caller-supplied state tensors are returned by identity:
+the suffix aliases the same buffers, which the next selective call updates in place.
+When only the public state suffix is supplied, temporary sorting scratch is allocated
+internally without replacing those state buffers.
+
+Use two phases for PyTorch selective COO. First make an eager bootstrap call with
+every flag true; it may allocate omitted topology and tile buffers and returns the
+reusable state. Every later eager call containing a false flag must supply all fixed
+topology and tile-state buffers, otherwise it raises `ValueError` before a kernel
+launch. `pair_offsets` starts at zero, is nondecreasing, and ends at the physical
+COO capacity; `pair_counts` stays within each segment.
+
+```python
+# Eager bootstrap: omit every persistent buffer.
+state = cluster_tile_neighbor_list(
+    positions, cutoff, cell, format="coo", return_state=True,
+    rebuild_flags=torch.ones(1, dtype=torch.bool, device=positions.device),
+)
+```
+
+The direct single-system `cluster_tile_neighbor_list` route also supports
+`torch.compile(fullgraph=True)` for this steady-state phase. Bootstrap and validate
+capacities eagerly, then compile a function that accepts fixed-shape buffers and
+`rebuild_flags`; compiled calls preserve false-flag data and rebuild true-flag data.
+False flags retain the fixed launch graph, so they do not promise zero preprocessing
+or kernel launches. The unified `neighbor_list(..., method="cluster_tile")` dispatch
+is unsupported under compilation because it cannot validate tensor-valued `pbc`
+without host synchronization. Compiled calls also omit
+host-synchronized offset and overflow diagnostics, so retain the eagerly validated
+capacities. The Warp COO query enforces physical output-buffer bounds as defense in
+depth, but mutated or malformed metadata remains unsupported. Batched cluster-tile
+fullgraph support is not provided.
+
+```python
+@torch.compile(fullgraph=True)
+def reuse(flags, neighbor_list, pair_offsets, pair_counts, shifts, num_tiles, row, col):
+    return cluster_tile_neighbor_list(
+        positions, cutoff, cell, format="coo", return_state=True,
+        rebuild_flags=flags, neighbor_list=neighbor_list, pair_offsets=pair_offsets,
+        pair_counts=pair_counts, neighbor_list_shifts=shifts, num_tiles=num_tiles,
+        tile_row_group=row, tile_col_group=col,
+    )
+```
+
 ### Dual Cutoff
 
 Compute two neighbor lists with different cutoffs simultaneously:

nvalchemiops/jax/neighbors/batch_cluster_tile.py

@@ -1409,7 +1409,9 @@ def batch_query_cluster_tile(
         with pair outputs.
     rebuild_flags : jax.Array, shape (S,), dtype=bool, optional
         Per-system selective rebuild flags. Requires ``tile_offsets``,
-        ``tile_counts``, and ``batch_idx``.
+        ``tile_counts``, and ``batch_idx``. For an empty batch, false-flag
+        segments retain their active pair and tile counts while true-flag
+        segments clear them; fixed-capacity arrays remain unchanged.
     tile_offsets : jax.Array, shape (S + 1,), dtype=int32, optional
         Per-system tile segment offsets. Required with ``rebuild_flags``.
     tile_counts : jax.Array, shape (S,), dtype=int32, optional
@@ -2029,7 +2031,10 @@ def batch_cluster_tile_neighbor_list(
     For ``format == "matrix"``:
         ``(neighbor_matrix, num_neighbors, neighbor_matrix_shifts)``, with
         optional ``(*, distances)`` and/or ``(*, vectors)`` appended when
-        ``return_distances`` / ``return_vectors`` is True.
+        ``return_distances`` / ``return_vectors`` is True. With
+        ``rebuild_flags``, append ``(tile_offsets, tile_counts, num_tiles,
+        tile_row_group, tile_col_group, tile_system)`` after one matrix triple,
+        or after both triples when ``cutoff2`` is provided.
     For ``format == "coo"``:
         ``(neighbor_list, neighbor_ptr, neighbor_list_shifts)`` in compact
         mode, or ``(neighbor_list, pair_offsets, pair_counts,
@@ -2162,6 +2167,29 @@ def batch_cluster_tile_neighbor_list(
         nn0 = jnp.empty(0, dtype=jnp.int32)
         ns0 = jnp.empty((0, int(max_neighbors), 3), dtype=jnp.int32)
         if format == "coo":
+            if selective:
+                empty_pair_counts = jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_pair_counts),
+                    previous_pair_counts,
+                )
+                empty_tile_counts = jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_tile_counts),
+                    previous_tile_counts,
+                )
+                return (
+                    previous_neighbor_list,
+                    pair_offsets,
+                    empty_pair_counts,
+                    previous_neighbor_list_shifts,
+                    tile_offsets,
+                    empty_tile_counts,
+                    previous_num_tiles,
+                    previous_tile_row_group,
+                    previous_tile_col_group,
+                    previous_tile_system,
+                )
             coo_base = (
                 jnp.empty((2, 0), dtype=jnp.int32),
                 jnp.zeros(1, dtype=jnp.int32),
@@ -2231,7 +2259,11 @@ def batch_cluster_tile_neighbor_list(
             return (
                 *matrix_out,
                 tile_offsets,
-                previous_tile_counts,
+                jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_tile_counts),
+                    previous_tile_counts,
+                ),
                 previous_num_tiles,
                 previous_tile_row_group,
                 previous_tile_col_group,

nvalchemiops/jax/neighbors/cluster_tile.py

@@ -1154,8 +1154,10 @@ def query_cluster_tile(
     neighbor_matrix_shifts2 : jax.Array, shape (natom, max_neighbors, 3), dtype=int32, optional
         Second shift buffer for dual-cutoff mode.
     rebuild_flags : jax.Array, shape (1,), dtype=bool, optional
-        When set, skips atoms flagged as unchanged; requires all
-        ``previous_*`` buffers to be passed.
+        When set, requires all ``previous_*`` buffers. For an empty system a
+        false flag preserves the previous active counts and tile state, while
+        a true flag clears the active pair and tile counts without rewriting
+        fixed-capacity topology storage.
     return_vectors : bool, optional
         If True, also fill and return per-pair displacement vectors.
     return_distances : bool, optional
@@ -1802,7 +1804,10 @@ def cluster_tile_neighbor_list(
     For ``format == "matrix"``:
         ``(neighbor_matrix, num_neighbors, neighbor_matrix_shifts)``, with
         optional ``(*, distances)`` and/or ``(*, vectors)`` appended when
-        ``return_distances`` / ``return_vectors`` is True.
+        ``return_distances`` / ``return_vectors`` is True. With
+        ``rebuild_flags``, append ``(num_tiles, tile_row_group,
+        tile_col_group)`` after one matrix triple, or after both triples when
+        ``cutoff2`` is provided.
     For ``format == "coo"``:
         ``(neighbor_list, neighbor_ptr, neighbor_list_shifts)`` in compact
         mode, or ``(neighbor_list, pair_offsets, pair_counts,
@@ -1920,6 +1925,26 @@ def cluster_tile_neighbor_list(
         nn0 = jnp.empty(0, dtype=jnp.int32)
         ns0 = jnp.empty((0, int(max_neighbors), 3), dtype=jnp.int32)
         if format == "coo":
+            if selective:
+                empty_pair_counts = jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_pair_counts),
+                    previous_pair_counts,
+                )
+                empty_num_tiles = jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_num_tiles),
+                    previous_num_tiles,
+                )
+                return (
+                    previous_neighbor_list,
+                    pair_offsets,
+                    empty_pair_counts,
+                    previous_neighbor_list_shifts,
+                    empty_num_tiles,
+                    previous_tile_row_group,
+                    previous_tile_col_group,
+                )
             coo_base = (
                 jnp.empty((2, 0), dtype=jnp.int32),
                 jnp.zeros(1, dtype=jnp.int32),
@@ -1981,7 +2006,11 @@ def cluster_tile_neighbor_list(
         if selective:
             return (
                 *matrix_out,
-                previous_num_tiles,
+                jnp.where(
+                    rebuild_flags,
+                    jnp.zeros_like(previous_num_tiles),
+                    previous_num_tiles,
+                ),
                 previous_tile_row_group,
                 previous_tile_col_group,
             )

nvalchemiops/jax/neighbors/naive_dual_cutoff.py

@@ -614,6 +614,79 @@ def naive_neighbor_list_dual_cutoff(
                 per_atom_cell_offsets,
                 launch_dims=(total_atoms,),
             )
+            if rebuild_flags is not None:
+                rf = rebuild_flags.flatten()[:1].astype(jnp.bool_)
+                num_neighbors1 = jnp.where(
+                    rf[0], jnp.zeros_like(num_neighbors1), num_neighbors1
+                )
+                num_neighbors2 = jnp.where(
+                    rf[0], jnp.zeros_like(num_neighbors2), num_neighbors2
+                )
+                (
+                    neighbor_matrix1,
+                    neighbor_matrix_shifts1,
+                    num_neighbors1,
+                    neighbor_matrix2,
+                    neighbor_matrix_shifts2,
+                    num_neighbors2,
+                ) = _jax_fill_pbc_selective(
+                    positions_wrapped,
+                    per_atom_cell_offsets,
+                    float(cutoff1 * cutoff1),
+                    float(cutoff2 * cutoff2),
+                    cell,
+                    shift_range_per_dimension,
+                    empty_num_shifts,
+                    empty_batch_idx,
+                    empty_batch_ptr,
+                    empty_target_indices,
+                    neighbor_matrix1,
+                    neighbor_matrix_shifts1,
+                    num_neighbors1,
+                    neighbor_matrix2,
+                    neighbor_matrix_shifts2,
+                    num_neighbors2,
+                    empty_vectors,
+                    empty_distances,
+                    empty_pair_params,
+                    empty_energies,
+                    empty_forces,
+                    rf,
+                    launch_dims=(1, max_shifts_per_system, total_atoms),
+                )
+            else:
+                (
+                    neighbor_matrix1,
+                    neighbor_matrix_shifts1,
+                    num_neighbors1,
+                    neighbor_matrix2,
+                    neighbor_matrix_shifts2,
+                    num_neighbors2,
+                ) = _jax_fill_pbc(
+                    positions_wrapped,
+                    per_atom_cell_offsets,
+                    float(cutoff1 * cutoff1),
+                    float(cutoff2 * cutoff2),
+                    cell,
+                    shift_range_per_dimension,
+                    empty_num_shifts,
+                    empty_batch_idx,
+                    empty_batch_ptr,
+                    empty_target_indices,
+                    neighbor_matrix1,
+                    neighbor_matrix_shifts1,
+                    num_neighbors1,
+                    neighbor_matrix2,
+                    neighbor_matrix_shifts2,
+                    num_neighbors2,
+                    empty_vectors,
+                    empty_distances,
+                    empty_pair_params,
+                    empty_energies,
+                    empty_forces,
+                    empty_rebuild_flags,
+                    launch_dims=(1, max_shifts_per_system, total_atoms),
+                )
         else:
             if rebuild_flags is not None:
                 rf = rebuild_flags.flatten()[:1].astype(jnp.bool_)