Emulator differentiability
==========================

The learned emulator differentiates its uniform sticky-sphere mechanics and
graph diffusion so their gradients compose with the GNS.

Frozen topology within each correction
--------------------------------------

Each learned update may be followed by ``mech_steps`` mechanics calls and
``diff_steps`` diffusion calls; zero selects the identity path for that
correction. ``dt_mech`` and ``dt_diff`` set their respective explicit step
sizes.

Pair discovery happens outside autodiff. The emulator rebuilds topology between
correction calls as the state evolves, then holds it fixed while that call is
differentiated. Gradients cover continuous force, flux, and state arithmetic;
neighbor topology remains a discrete forward input. Learned graphs use a host
:class:`scipy.spatial.cKDTree`; prescribed emulator corrections use a Warp
HashGrid. Candidate and output work depends on local density and can be
quadratic for dense states.

Torch Tape and JAX custom VJP
-----------------------------

:mod:`waxmorph.torch.warp_autograd` wraps each mechanics or diffusion call in a
:class:`torch.autograd.Function` with a fresh :class:`warp.Tape`. Backward seeds
the output adjoint and replays the recorded emulator kernels. The mechanics
bridge exposes position gradients; the diffusion bridge exposes concentration
gradients.

:mod:`waxmorph.jax.warp_autograd` uses custom VJPs on CUDA, backed by dedicated
backward-enabled per-pair Warp FFI kernels. JAX owns endpoint gathers,
padding masks, scatter-adds, and state updates. Its mechanics bridge
differentiates positions and radii on the frozen pairs; diffusion differentiates
concentrations. Adding a correction therefore requires backend-specific bridge
work.

Simulator derivative implementations
====================================

The forward simulator carries its own derivative code, separate from the
emulator bridges above. ``mech_step_sticky`` uses explicit pair derivatives.
``mech_step_sticky_implicit`` uses local :func:`warp.grad` evaluation for
polarity, thickness, mesenchymal alignment, and optional WNT terms; its
soft-sphere derivatives remain explicit. Both paths use atomic reductions, so
floating-point reduction order may vary on parallel devices.