Add parallel MALA sampler (Phase 3)

Implement parallelized Metropolis-Adjusted Langevin Algorithm using DEER.

MALA components in src/parallel/mala.jl:
- MALARandomInputs: Pre-sampled (ξ, u) pairs for proposals and accept-reject
- MALAConfig: Configuration struct with step size, log density, gradient
- mala_proposal(): Langevin proposal x̃ = x + ε∇log p(x) + √(2ε)ξ
- mala_log_acceptance_ratio(): MH ratio with forward/backward densities
- soft_gate(): Differentiable accept-reject using sigmoid + straight-through
- mala_transition(): Complete MALA step combining proposal and accept-reject
- parallel_mala(): Run full MALA chain in parallel via DEER
- sequential_mala(): Reference implementation for testing

The stop-gradient trick enables computing Jacobians through the
non-differentiable accept-reject step by using a soft sigmoid gate
in the backward pass while keeping hard decisions in the forward pass.

Tests: 31 new tests covering proposal mechanics, sequential/parallel
equivalence, various target distributions, acceptance behavior, and
sample statistics.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: xukai92

docs/parallel_mcmc_implementation_plan.md

@@ -70,21 +70,23 @@ This document tracks the implementation of DEER (Doubly Efficient Estimation via
 
 ---
 
-### Phase 3: MALA Integration ⬜
+### Phase 3: MALA Integration ✅
 > Parallelized Metropolis-Adjusted Langevin Algorithm
 
-- [ ] **3.1 MALA Transition Function**
-  - [ ] Proposal: x̃ = x + ε∇log p(x) + √(2ε)ξ
-  - [ ] Acceptance ratio computation
-  - [ ] Stop-gradient trick for differentiable accept-reject
-  - [ ] Soft gating: g = σ(log α - log u) with straight-through estimator
+- [x] **3.1 MALA Transition Function** ✅
+  - [x] Proposal: x̃ = x + ε∇log p(x) + √(2ε)ξ
+  - [x] Acceptance ratio computation (forward and backward proposal densities)
+  - [x] Stop-gradient trick for differentiable accept-reject
+  - [x] Soft gating with sigmoid and straight-through estimator
 
-- [ ] **3.2 Parallel MALA Sampler**
-  - [ ] Pre-sample all random inputs (ξ for proposals, u for accept-reject)
-  - [ ] Integrate with DEER framework
-  - [ ] Return full chain from parallel computation
+- [x] **3.2 Parallel MALA Sampler** ✅
+  - [x] MALARandomInputs type for pre-sampled (ξ, u) pairs
+  - [x] sample_mala_inputs() for batch sampling
+  - [x] parallel_mala() integrating with DEER framework
+  - [x] sequential_mala() for reference/testing
+  - [x] Convenience API with automatic input sampling
 
-- [ ] **3.3 MALA-specific Optimizations**
+- [ ] **3.3 MALA-specific Optimizations** (deferred)
   - [ ] Preconditioning with Hessian eigendecomposition (optional)
 
 ---
@@ -197,14 +199,14 @@ src/
 │   ├── scan.jl              # Parallel scan implementations ✅
 │   ├── jacobian.jl          # Jacobian computation utilities ✅
 │   ├── deer.jl              # Core DEER algorithm ✅
-│   ├── mala.jl              # Parallel MALA (TODO)
+│   ├── mala.jl              # Parallel MALA ✅
 │   └── hmc.jl               # Parallel HMC (TODO)
 test/
 ├── parallel/
 │   ├── test_scan.jl         # ✅ 141 tests passing
 │   ├── test_jacobian.jl     # ✅ 57 tests passing
 │   ├── test_deer.jl         # ✅ 67 tests passing
-│   ├── test_mala.jl         # (TODO)
+│   ├── test_mala.jl         # ✅ 31 tests passing
 │   └── test_hmc.jl          # (TODO)
 ```
 
@@ -249,6 +251,10 @@ test/
 | 2026-01-20 | 2 | DEER algorithm core | ✅ | Newton iteration, Full/Quasi/Stochastic DEER |
 | 2026-01-20 | 7.1 | DEER algorithm tests | ✅ | 67 tests passing |
 | 2026-01-20 | 2 | **Phase 2 Complete** | ✅ | Core DEER algorithm working |
+| 2026-01-20 | 3 | MALA transition function | ✅ | Proposal, acceptance, soft gating |
+| 2026-01-20 | 3 | Parallel MALA sampler | ✅ | Integrated with DEER framework |
+| 2026-01-20 | 7.1 | MALA tests | ✅ | 31 tests passing |
+| 2026-01-20 | 3 | **Phase 3 Complete** | ✅ | Parallel MALA working |
 
 ---
 

src/parallel/Parallel.jl

@@ -47,6 +47,9 @@ include("jacobian.jl")
 # Core DEER algorithm
 include("deer.jl")
 
+# Parallel MALA
+include("mala.jl")
+
 # Export types
 export AbstractParallelMethod, FullDEER, QuasiDEER, StochasticQuasiDEER, BlockQuasiDEER
 
@@ -76,4 +79,9 @@ export hessian_diagonal, batch_hessian_diagonals
 export DEERResult
 export deer, deer_with_settings, sequential_mcmc
 
+# Export MALA
+export MALARandomInputs, MALAConfig
+export sample_mala_inputs, mala_proposal, mala_transition
+export parallel_mala, sequential_mala
+
 end # module