Recorded-Mesh Neural ODE Backpropagation

1. Integration point: replace the differentiable solver inside a neural ODE or probability-flow diffusion training loop. The vector field is f_theta(t,y,c), where c is conditioning and y contains a batch or ensemble of initial states. In a mesh-recording pass, run an embedded adaptive RK solver with tolerances tau_abs and tau_rel, store each trajectory's accepted step sizes h_i, and stop gradients through this pass. Pad trajectories to S_max and attach an active mask; padded steps use h=0 or an identity map. 2. Replay pseudocode: record = stop_gradient(adaptive_solve(f_theta,y0,t0,t1)); h,mask = pad(record.steps); y=y0; t=t0; for i in range(S_max): y_old=y; y=rk_step(f_theta,y,t,h[:,i],theta); y=where(mask[:,i],y,y_old); t=t+h[:,i]; loss=objective(y,target); backward(loss). Implement rk_step with the displayed stage equations and fixed coefficients, using JAX lax.scan, XLA fusion, or a custom CUDA kernel. 3. The mathematics determines the accepted-step record, fixed RK map, and discrete-adjoint semantics. Empirically monitor replay-versus-adaptive state error, gradient cosine similarity, and trajectories exceeding S_max. Refresh the mesh every K updates or when a validation adaptive solve reports RMS error above twice the target tolerance. 4. First experiment: train a 2-layer MLP neural ODE on a 2D spiral or downsampled MNIST, comparing Diffrax adaptive backpropagation, fixed-step RK4, and recorded-mesh replay at matched forward error and batch size. Measure wall-clock time per update, peak memory, loss at equal GPU time, terminal-state error, and gradient cosine similarity. Success means at least 2x lower gradient-step time or lower memory at equal final accuracy without unstable gradients.