ParaRNN [from Apple Research] (github.com)

Apple Research released ParaRNN, an open-source package and accompanying paper that enables large speedups by applying nonlinear RNNs in parallel along sequence length. ParaRNN replaces the classic step-by-step hidden-state update with a Newton-linearization plus parallel-reduction solver, letting RNN recurrence be evaluated concurrently across timesteps. The library ships reference diagonalized GRU/LSTM cells, a PyTorch-friendly API that autogenerates Jacobians via autograd, and high-performance CUDA kernels—making it practical to prototype in pure PyTorch and scale to GPU-accelerated production without reimplementing whole models. Technically, ParaRNN assembles the Newton linearized system for arbitrary RNN cells and solves it with parallel-reduction algorithms optimized for diagonal and block-diagonal Jacobian structures. It exposes four application modes—from classical sequential to fully fused CUDA kernels—so users can trade ease-of-use for peak performance (parallel_CUDA uses PyTorch for assembly and custom CUDA for reduction; parallel_FUSED requires CUDA implementations of recurrence and Jacobians for the fastest end-to-end kernel). Installation requires Python 3.9+, PyTorch, CUDA and a C++ toolchain; note numerical error grows with sequence length and machine precision and users must ensure Newton stability (bounded Jacobians) when designing cells. ParaRNN makes experimenting with new RNN designs fast while unlocking parallel training performance previously unavailable for long-sequence RNN workloads.