PCAX — Predictive Coding in JAX

PCAX is a research-grade Predictive Coding(PC) framework built in JAX. The entire network state is a single 1D array and connections are typed block-sparse weights, which makes jit, grad, vmap, and scan work over the whole model for free.

A key design point: users define only the top-down (generative) edges. The bottom-up inference pathway — error backpropagation through Wᵀ — is derived automatically via jax.grad(pc_energy).

This repository contains two things:

• pcax/ — the reusable library (this is the package you install and build on).
• papers/eccv2026-wfactpc/ — the reproduction bundle for our ECCV 2026 paper (WF-Act-PC). See papers/eccv2026-wfactpc/REPRODUCE.md.

Install

# Library only (jax, jaxlib, optax)
pip install -e .

# Library + extras needed to run the paper experiments (torch, torchvision, ...)
pip install -e ".[experiments]"

Requires Python ≥ 3.11 and a JAX build matching your accelerator (CUDA 12 for GPU).

Matmul precision: deep PC networks (depth > ~10) with μPC are sensitive to TF32 rounding on Ampere GPUs. Set jax.config.update("jax_default_matmul_precision", "float32") before any JAX op.

Quickstart

import jax, optax
from pcax import (
    GraphSpec, compile_graph, init_weights, make_prediction_scalings,
    WeightStore, relaxation_step, pc_energy,
)

jax.config.update("jax_default_matmul_precision", "float32")

# 1. Define the generative model (top-down edges only)
spec = GraphSpec()
spec.add_node("input",  size=784, activation="linear", clamped=True)
spec.add_node("hidden", size=128, activation="relu")
spec.add_node("output", size=10,  activation="linear", clamped=True)
spec.add_dense_edge("input",  "hidden")
spec.add_dense_edge("hidden", "output")

# 2. Compile + choose a parameterisation (SP / muPC)
system  = compile_graph(spec)
weights = init_weights(spec, jax.random.PRNGKey(0))
system  = system._replace(prediction_scalings=make_prediction_scalings(system, "mupc"))

# 3. Inference (relax activities) and learning (update weights at equilibrium)
@jax.jit
def infer_step(params, states):
    return relaxation_step(WeightStore(params=params), states, system,
                           lr=0.5, loss_type="mse", output_node_name="output")

See pcax/README.md for the full architecture, data structures, block types, dynamics, and extension points.

Citation

If you use PCAX, please cite the WF-Act-PC paper (ECCV 2026).

License

Licensed under the Apache License, Version 2.0. See LICENSE for details.