Open source
Contributions to scientific and machine-learning software.
This page lists my merged contributions to upstream projects and the libraries I develop. More code is on GitHub.
Upstream projects
Contributions.
-
Race-free GPU Gauss–Seidel smoothing at cut cells in MLEBABecLap
Merged · approved by Weiqun Zhang
In the embedded-boundary linear solver of AMReX (MLEBABecLap), the red–black Gauss–Seidel smoother reads diagonal neighbours of a cut cell during face-centroid flux interpolation and in the Dirichlet boundary gradient. These neighbours have the same red–black parity as the cell being updated. On a GPU, a thread can therefore read a value that another thread of the same launch is writing, and the result depends on thread scheduling. I reported the race in issue #6092.
The pull request adds an opt-in deterministic mode, enabled with LPInfo::setDeterministic(true), in which each half-sweep is split into 2^(D−1) colours so that no two cells in a 3 × 3 × 3 neighbourhood share a colour. The default path is unchanged. On an NVIDIA RTX A6000, repeated multigrid solves in deterministic mode were bitwise identical in every tested configuration, and the solve time changed by −1 to +37 per cent.
-
LoRA weight synchronization in the FSDP workers
Merged
With LoRA enabled and the base weights already synchronized, the FSDP workers in VERL passed the per-tensor parameters to the bucketing routine as a dictionary. The routine expects an iterator of (name, tensor) pairs, so weight synchronization failed with a ValueError. The fix passes params.items().
Libraries
Libraries I develop.
-
JAX-DIPS
A differentiable solver for three-dimensional interfacial PDEs, written in JAX. It runs on GPUs, TPUs and CPUs. The solver trains compact neural networks to solve elliptic problems with jump conditions at interfaces, using neural bootstrapping of finite discretization methods.
-
Stencil plugin
The companion plugin of the Stencil journal for coding agents such as Claude Code and Codex. Each essay contributes skills, worked examples and verifiable reports for physics simulation, and the plugin collects them in one versioned package in which each contribution is attributed to its author.