How does Parallax compare to Flax for JAX?

Parallax offers a PyTorch-like module API with immutable trees and explicit state management, while Flax's Linen API is more mature with a larger ecosystem. Choose Parallax for pure functional design experiments, but Flax for production-ready projects.

How to implement dropout with random number generation in Parallax?

RNG is handled by attaching seeds to the tree; use `self.rng` in forward passes and `.new_state()` to update states, as demonstrated in the Dropout example where RNG is used once or split for stochastic operations.

Is Parallax stable enough for long-term research projects?

No, it's explicitly a prototype, so APIs may change, and it lacks the stability and support needed for critical or long-term research, making it better for short-term experiments or prototyping.

Can I convert my PyTorch models to Parallax easily?

Partially; the module syntax is similar, but you must rewrite initialization to use `ParamInit` and adapt to immutable state management, which requires a functional programming mindset.

What are the performance costs of immutable trees in Parallax?

Immutable trees can incur memory overhead due to copying on state updates, but JAX's optimizations may help. The trade-off is improved debuggability and alignment with JAX's functional transformations.

How do I share parameters between modules in Parallax?

Parameter sharing is explicit through the tree structure; define shared modules as attributes and reference them, ensuring immutability, as seen in the recurrent use of `dense2` in the BinaryNetwork example.

Parallax — Immutable JAX Module System

What is Parallax?

Parallax is a prototype module system for JAX that introduces immutable parameter trees, eliminating the need for imperative state tracking and initialization patterns common in frameworks like PyTorch. It enables pure functional neural network definitions while maintaining familiar module-based APIs, making it easier to reason about parameter sharing, random number generation, and gradient updates in JAX-based machine learning workflows.

Target Audience

Machine learning researchers and engineers using JAX who want to define neural networks with object-oriented, PyTorch-like modules while adhering to JAX's functional programming principles.

Value Proposition

Developers choose Parallax over alternatives because it provides immutable, declarative module definitions that eliminate hidden global state and imperative patterns, making parameter sharing, RNG handling, and gradient updates more explicit and easier to manage within JAX's functional ecosystem.

Overview

Key Features

Immutable Parameter Trees — All modules and parameters are immutable trees, replacing imperative coding patterns with declarative structures.
Explicit Parameter Declaration — Parameters are explicitly declared as class attributes using Parameter and ParamInit, making module dependencies clear.
Pure Functional Forward Passes — Forward methods operate on immutable state, with random number generation handled via tree-attached RNG seeds.
Declarative Initialization — Parameter initialization is separated from module definition via an .initialized() method that binds random seeds.
State Threading — Training/evaluation modes and RNG state are managed by threading through the parameter tree with .new_state().
JAX-Compatible Gradients — Gradient computation and updates work directly on the immutable tree structure using JAX transformations.

Philosophy

Parallax embraces pure functional programming principles by making all module state immutable and explicitly managed, avoiding hidden global state while maintaining the usability of object-oriented neural network definitions.

Parallax

What is Parallax?

Overview

Key Features

Philosophy

Related Projects

Found a gem we're missing?

Use Cases

Best For

Not Ideal For

Pros & Cons

Pros

Cons

Open Source Alternative To

Frequently Asked Questions