🚀[FEA]: JAX neighbor lists require a static `cutoff` (can't be traced)

[n-gao]

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Critical (currently preventing usage)

Please provide a clear description of problem you would like to solve.

batch_cell_list / batch_naive_neighbor_list coerce cutoff to a Python
float, so building a neighbor list inside jax.jit with a traced cutoff
(e.g. from parameters) fails.

Reproducer

import jax, jax.numpy as jnp
from nvalchemiops.jax.neighbors import batch_cell_list

pos = jnp.zeros((3, 3), jnp.float32)
cell = (jnp.eye(3, dtype=jnp.float32) * 10.0)[None]
batch_ptr = jnp.array([0, 3], jnp.int32)

@jax.jit
def build(cutoff):
    return batch_cell_list(pos, cutoff, cell=cell, pbc=jnp.ones((1, 3), bool),
                           batch_ptr=batch_ptr, max_neighbors=8, max_total_cells=64)

build(jnp.asarray(5.0, jnp.float32))
# jax.errors.ConcretizationTypeError: ... traced array with shape float32[]
#   The problem arose with the `float` function.

Versions: nvalchemiops 0.3.1, warp-lang 1.13.0, jax 0.10.0 (CUDA).

Describe any alternatives you have considered

No response

[zubatyuk]

Hi Nicholas.
Is cutoff genuinely data-dependent/trainable inside jax.jit, or could it be made static with static_argnames? I’m asking because Warp’s JAX scalar arguments are static today. Supporting dynamic cutoffs would require separate array-backed kernel variants, while static_argnames would recompile for each distinct cutoff value. Is that recompilation cost or cutoff variability the reason static cutoff is not an option here?

[n-gao]

In https://github.com/cusp-ai-oss/kUPS, we generally try to keep as much dynamic as possible for maximum compatibility and cutoffs are generally stored in Jax arrays. We would need to refactor a few code paths to accommodate static cutoff radii. Nothing impossible but we'd generally prefer it to be traceable if possible.

One potential application would be:
Han et al. "Smooth Dynamic Cutoffs for Machine Learning Interatomic Potentials"