waxmorph.jax.warp_autograd

JAX/Warp custom-VJP bridge for frozen-pair physics on CUDA.

Backward-enabled Warp FFI differentiates dedicated per-pair kernels; JAX owns endpoint gathers, masking, scatter aggregation, and state updates. pair_i and pair_j are constructed outside autodiff, so gradients treat topology as fixed. PyTorch records its emulator kernels on a warp.Tape; JAX uses per-pair FFI kernel objects.

Top-level JAX APIs are preferred; a lazy legacy fallback supports the declared Warp 1.10 minimum.

Functions

warp_diffusion_step(c, pair_i, pair_j, ...)	Apply one differentiable graph-Laplacian signaling-molecule diffusion step.
warp_mech_step(X, R, pair_i, pair_j, dt, *)	Apply one differentiable sticky-sphere mechanics step.