Context
From Gemini code review on PR #23: the shared_params=False path in _compute_transient_fd_gradients (circulax/solvers/adjoint.py:343-363) uses a nested Python loop over n_devices × n_params, making eager osdi_residual_eval FFI calls on each iteration.
For larger circuits this O(params × devices) eager dispatch becomes a significant bottleneck in transient sensitivity analysis.
Current code
for pi, (_pname, pcol) in enumerate(zip(param_names, param_cols, strict=True)):
for i in range(n_devices):
params_perturbed = params_np.copy()
# ... perturb one element ...
cur_pert, chg_pert_cur, _ = osdi_residual_eval(mid, v_all_cur, params_jax_pert, group.states)
# ... two FFI calls per (param, device) pair ...Proposed improvement
Batch perturbation columns into fewer FFI calls — e.g. perturb all devices for one parameter in a single batched osdi_residual_eval call, or restructure to perturb across the full (N, num_params) matrix and evaluate once.
This requires changes to the bosdi FFI boundary to support batched perturbation evaluation.
Priority
Medium — correctness is fine, this is a performance optimization for large OSDI circuits.
Context
From Gemini code review on PR #23: the
shared_params=Falsepath in_compute_transient_fd_gradients(circulax/solvers/adjoint.py:343-363) uses a nested Python loop overn_devices × n_params, making eagerosdi_residual_evalFFI calls on each iteration.For larger circuits this O(params × devices) eager dispatch becomes a significant bottleneck in transient sensitivity analysis.
Current code
Proposed improvement
Batch perturbation columns into fewer FFI calls — e.g. perturb all devices for one parameter in a single batched
osdi_residual_evalcall, or restructure to perturb across the full(N, num_params)matrix and evaluate once.This requires changes to the bosdi FFI boundary to support batched perturbation evaluation.
Priority
Medium — correctness is fine, this is a performance optimization for large OSDI circuits.