Neural Solver Usage

The neural_solver package is a GPU-native, mesh-free companion to the Rust solver crate. It implements neural methods for backward stochastic differential equations (BSDE) and Hamilton-Jacobi-Bellman (HJB) equations in JAX.

This page covers the package itself. The theory is on the theory pages and measured results are in the empirical results section.

The Rust solver crate is the rigorous reference: it solves the same problems with finite-difference policy iteration, least-squares Monte Carlo BSDE, and analytical solutions, all validated against closed forms. neural_solver adds a differentiable, device-resident path that targets problems where a fixed grid becomes impractical (high state dimension).

Why a separate package

neural_solver is a Python package, not a Cargo workspace member. JAX targets Linux with NVIDIA (CUDA) or AMD (ROCm) accelerators, independent of the Rust toolchain. The two codebases are coupled only through committed reference values, not a live call graph: the Rust solver emits exact values, and the JAX tests assert agreement against them.

Design principles

• Device-resident training. The forward simulation, backward driver, loss, and optimizer step all run on-device; the host reads only scalars.
• Differentiate through the problem. Value and control networks are differentiated with JAX autodiff, not finite differences.
• Validated against the Rust solver. Every neural method is checked against exact or committed reference values.

Layout

• sde.py: differentiable forward SDE simulation via jax.lax.scan. The diffusion-only euler_maruyama and the jump-aware euler_maruyama_jump (single compound Poisson source) are implemented.
• bsde.py: deep BSDE solver. train_deep_bsde handles diffusions (full Z-process method); train_deep_bsde_jump adds a learned Gamma process and jump sampling.
• networks.py: tanh/softplus/relu Mlp used by the deep BSDE and DGM methods.
• loss.py: terminal-matching loss.
• dgm.py: deep HJB (DGM) helpers: the correct spatial Hessian diagonal and the diffusion term. The full DGM residual is not yet implemented.
• models/: one file per reference problem (Black-Scholes, Merton, Merton with deterministic jumps, Merton with log-normal jumps, LQ, jump-LQ, Avellaneda-Stoikov). jump_lq.py provides the Poisson-jump LQ regulator with its Riccati closed form and a decoupled BSDE problem for validating the Gamma path. merton_jump.py provides the deterministic-jump deep BSDE target; merton_jump_lognormal.py is reference-only because the deep BSDE loop does not yet sample random jump sizes.

Runnable Python examples live in neural_solver/README.md; they are not embedded here because the mdbook examples must remain valid Rust.

Installation

See neural_solver/README.md for the full installation matrix. In short:

• CPU only: python -m pip install -e .
• NVIDIA (CUDA 12): python -m pip install -e ".[cuda]"
• AMD (ROCm 7.14, Linux): install rocm[libraries,device-gfx1201] and jax_rocm7_plugin/jax_rocm7_pjrt from the AMD wheel index, then jax==0.10.0/jaxlib==0.10.0, then python -m pip install -e ..

JAX-on-ROCm is Linux-only. Native Windows and WSL2 are not currently distributed for JAX, even though the ROCm runtime itself supports Windows for the RX 9070 (gfx1201).

An unattended Ubuntu Desktop 24.04 installer that provisions amdgpu, ROCm 7.14.0, GitHub Desktop, and Zed is provided in neural_solver/autoinstall/; see its README.md.

Reference agreement

Neural methods are validated against exact reference values produced by the Rust solver crate. The first benchmark in the ladder is pure agreement, with no training: the JAX closed forms for Merton, Merton with jumps, LQ, and Avellaneda-Stoikov are compared against the Rust exact values. Black-Scholes is validated against its own JAX closed form; the Rust generator emits no Black-Scholes entry.

Reference schema

The committed fixture neural_solver/tests/data/reference.json has five entries, each produced by solver/examples/reference/emit_reference.rs:

The scalar LQ entry uses the Rust LqRegulator<1,1> with a=-0.5, b=1, c=0.2, q=1, q_terminal=1, r=1, horizon=1. The Avellaneda-Stoikov entry uses AvellanedaExact with gamma=0.5, sigma=0.5, kappa=1.5, a=140, terminal_time=1, q_max=10. The theta values are gauge-relative (see the issue recorded in docs/src/project/issues.md); the spreads are gauge-invariant and are the physical outputs.

Generating reference samples

Reference values are produced by the Rust example solver/examples/reference/emit_reference.rs, which prints a single JSON object to stdout. To regenerate the committed fixture, run from the repository root:

cargo run -p solver --example emit_reference > neural_solver/tests/data/reference.json

Then verify the JAX closed forms still agree:

cd neural_solver && python -m pytest tests/test_reference_match.py

The fixture neural_solver/tests/data/reference.json is checked in so the JAX tests run without requiring a Rust toolchain. Regenerate it whenever the reference problems, parameters, or exact solutions change.

See the neural_solver workstream for the stages and their definitions of done.

Usage

Runnable usage examples live in neural_solver/README.md; they are not embedded here because the mdbook examples must remain valid Rust. The main entry points are:

• neural_solver.sde.euler_maruyama for differentiable forward simulation.
• neural_solver.bsde.train_deep_bsde for the full deep BSDE training loop.
• neural_solver.dgm.hessian_diagonal and neural_solver.dgm.diffusion_term for the DGM diffusion helpers.
• neural_solver.models for closed-form reference values.

See neural_solver/README.md for installation and a runnable example.