Skip to content

Integrated Information Theory accelerated in JAX โšก

PyPI tests docs license python

iitx computes integrated information โ€” ฯ†_s, big ฮฆ, and the full cause-effect structure of IIT 3.0 and 4.0 โ€” for any finite discrete dynamical system given as a transition probability matrix. It reproduces the reference implementation PyPhi's numbers on the canonical examples, and does four things PyPhi cannot:

  • ๐Ÿš€ Parallelize within a system โ€” candidate systems, mechanisms, purviews, partitions, and cuts are batched tensor work on CPU, GPU, or TPU, not Python loops.
  • ๐Ÿ“ฆ Batch across systems โ€” every analysis is a pure function of (system, state) that composes with jax.vmap: ten thousand systems are one call. On a laptop CPU, that is ~100,000 exact ฯ†_s evaluations per second at n = 3 โ€” enough to evaluate every deterministic 3-unit system, exhaustively, in three minutes.
  • ๐Ÿง  Differentiate โ€” the IIT 4.0 pipeline is differentiable with respect to the system's dynamics, almost everywhere and with exact subgradients: gradient ascent on integrated information itself, a first.
  • ๐Ÿงฉ Treat the measure as a choice โ€” IIT 3.0, IIT 4.0 (2023), and IIT 4.0 (2026) are pluggable measures over one substrate, so the same system can be scored under competing formalizations in one line.

โœจ Overview

  • General. Binary or multi-valued units, uniform or heterogeneous alphabets, deterministic or stochastic dynamics. Coarse-graining and black-boxing are pure System โ†’ System transforms, so macro-level systems feed the unchanged pipeline.
  • Faithful. PyPhi (pinned) is the frozen test oracle: the regression suite compares values against generated golden fixtures for both theory versions, including non-binary systems, and every deliberate divergence is a documented finding (docs/notes/oracle-findings.md), never a silent patch.
  • Honest about exactness. Everything is exact; the IIT 3.0 earth mover's distances are host linear programs (so 3.0 jits and vmaps but refuses grad loudly), while the IIT 4.0 intrinsic-difference pipeline is pure tensor algebra end to end. Named, clearly-labeled relaxations (iitx.relax) exist for optimization landscapes.

๐Ÿ“ฆ Installation

pip install iitx

Requires Python โ‰ฅ 3.14. Dependencies: jax, numpy, and POT (the exact earth-mover's-distance backend for IIT 3.0).

๐Ÿš€ Getting started

import jax
import jax.numpy as jnp

from iitx.measures import iit4
from iitx.system import System

# The standard 3-unit example (OR, COPY, XOR), as a state-by-node TPM.
tpm = jnp.asarray(
    [
        [0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0],
        [1.0, 0.0, 1.0],
        [1.0, 0.0, 0.0],
        [1.0, 1.0, 0.0],
        [1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0],
        [1.0, 1.0, 0.0],
    ]
)
system = System.from_state_by_node(tpm)
state = jnp.asarray([1, 0, 0])

# System integrated information and the full ฮฆ-structure.
analysis = iit4.system_phi(system, state)
structure = iit4.phi_structure(system, state)
print(analysis.phi, structure.big_phi)  # 0.415... and 1.0

# Batch over all 8 states of the system in one call.
states = jnp.stack([jnp.asarray([s >> i & 1 for i in range(3)]) for s in range(8)])
batched = jax.vmap(iit4.system_phi, in_axes=(None, 0))(system, states)

# Differentiate ฯ†_s with respect to the dynamics.
gradient = jax.grad(
    lambda tpm: iit4.system_phi(System(tpm, system.shape, system.cm), state).signed_phi
)(system.tpm)

The full documentation lives at maxencefaldor.github.io/iitx.

๐Ÿ““ Examples

Notebook What it shows
Getting started Systems, ฯ†_s, ฮฆ-structures, IIT 3.0, batching, non-binary units, macro levels
Ascending ฮฆ Gradient ascent on integrated information itself โ€” to our knowledge a first โ€” with structure emerging along the way

๐Ÿ”ฌ Research

iitx is also a lab. The experiments/ directory holds a running research program on the gradients and landscapes of ฮฆ โ€” executed notebooks with pre-registered predictions, in reading order. Highlights so far: an audit showing exact ฯ†_s gradients match finite differences to 10โปยนโฐ (refuting the field's non-differentiability claim), the first exhaustive ฮฆ landscape (all 16.7M deterministic 3-unit systems, oracle-verified), the discovery that the 2023-formalism maximum is a frozen state reachable only from itself and grows as n(nโˆ’1), and that the 2026 revision zeroes the entire deterministic universe. A paper draft distilling the story lives in paper/.

๐Ÿ“ Design

The design documents are part of the repository: docs/design.md (the core abstractions and why), docs/generality.md (the PyPhi-vs-iitx capability table), and docs/glossary.md (one vocabulary, used identically in code and papers). Everything learned the hard way lives in docs/notes/.

๐Ÿ“– Citing iitx

If you use iitx in your research, please cite it (see also CITATION.cff):

@software{faldor2026iitx,
    author = {Faldor, Maxence},
    title = {iitx: Integrated Information Theory accelerated in JAX},
    url = {https://github.com/maxencefaldor/iitx},
    year = {2026},
}