blocktrix 0.0.5

A JAX library for solving block tri-diagonal matrix systems

blocktrix

A JAX library for efficiently solving block tri-diagonal matrix systems on GPUs.

Author: Philip Mocz (@pmocz)

Implements both a block Thomas (serial) solver and a block cyclic reduction B-cyclic (parallel) solver.

⚠️ Warning: Work in Progress ⚠️

This library is still under active development and is not guaranteed to work at this point XXX.

Installation

pip install blocktrix

Usage

import jax
from blocktrix import solve_block_tridiagonal_bcyclic, random_block_tridiagonal

# Generate a random test system
key = jax.random.PRNGKey(42)
n_blocks, block_size = 8, 4

lower, diag, upper, rhs = random_block_tridiagonal(key, n_blocks, block_size)

# Solve the system
x = solve_block_tridiagonal_bcyclic(n_blocks, lower, diag, upper, rhs)

It's fast!

Links

• Code repository
• Documentation

Cite this repository

If you use this software, please cite it as below.

@software{Mocz_Blocktrix_2026,
   author = {Mocz, Philip},
      doi = {https://doi.org/10.5281/zenodo.18226420},
    month = jan,
    title = {{Blocktrix}},
      url = {https://github.com/pmocz/blocktrix},
  version = {0.0.4},
     year = {2026}
}