splineax

Sparse linear operators and direct solvers for Lineax.

splineax lets you keep a linear system in its native sparse storage (jax.experimental.sparse.BCOO / BCSR) and solve it with a sparse direct solver that plugs straight into lineax.linear_solve. It also interfaces with asdex for calculating sparse Jacobians and using them as operators.

Operators: BCOOLinearOperator, BCSRLinearOperator, SparseJacobianLinearOperator.
Solvers: Spsolve (any backend), KLU (CPU-only, SuiteSparse, factorization reuse), and AutoSparseLinearSolver (picks one based on the platform).
A SparseLinearSolver protocol for separating factorization from solving.

Installation

pip install splineax

Example

Solve a 10000 x 10000 system. As a dense matrix it would need 10^8 entries, but kept sparse it has only ~3 x 10^4 nonzeros, and the solver never materialises the dense form.

import jax.numpy as jnp
import lineax as lx
import numpy as np
from jax.experimental.sparse import BCOO

import splineax as splx

n = 10000
np.random.seed(0)

# A large, randomly sparse matrix with a heavy diagonal (so it is invertible).
diagonal_indices = np.stack([np.arange(n), np.arange(n)], axis=1)
off_diagonal_indices = np.unique(np.random.randint(0, n, size=(2 * n, 2)), axis=0)
indices = jnp.concatenate([diagonal_indices, off_diagonal_indices])
values = jnp.concatenate(
    [
        np.full(n, float(n)),
        np.random.uniform(low=-1, high=1, size=off_diagonal_indices.shape[0]),
    ]
)
matrix = BCOO((values, indices), shape=(n, n)).sum_duplicates()

operator = splx.BCOOLinearOperator(matrix)
vectors = [jnp.ones(n), jnp.arange(n) % 2]
solver = splx.AutoSparseLinearSolver()

# Calculate factorization once...
with solver.factorize(operator) as factorized_state:
    # ...and reuse for multiple solves.
    solution = lx.linear_solve(
        operator, vectors[0], solver=solver, state=factorized_state
    )
    assert jnp.allclose(matrix @ solution.value, vectors[0], atol=1e-4)

    solution = lx.linear_solve(
        operator, vectors[1], solver=solver, state=factorized_state
    )
    assert jnp.allclose(matrix @ solution.value, vectors[1], atol=1e-4)

Documentation

Build the docs locally with uv run mkdocs serve, or view the user guide and API reference here.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

splineax-0.2.2.tar.gz (30.0 kB view details)

Uploaded Jul 26, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

splineax-0.2.2-py3-none-any.whl (40.9 kB view details)

Uploaded Jul 26, 2026 Python 3

File details

Details for the file splineax-0.2.2.tar.gz.

File metadata

Download URL: splineax-0.2.2.tar.gz
Upload date: Jul 26, 2026
Size: 30.0 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/6.1.0 CPython/3.13.14

File hashes

Hashes for splineax-0.2.2.tar.gz
Algorithm	Hash digest
SHA256	`f551ed828774640792beb6f2bc60380d72472dbb82eb6e4cf8e9a92a639845e0`
MD5	`8830391320e723460ef15e780214d026`
BLAKE2b-256	`9ce96e919e2850d0e63b256e410d155a0357da23c6f1333047f017d38ed05b7e`

See more details on using hashes here.

Provenance

The following attestation bundles were made for splineax-0.2.2.tar.gz:

Publisher: release.yml on nardi/splineax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: splineax-0.2.2.tar.gz
- Subject digest: f551ed828774640792beb6f2bc60380d72472dbb82eb6e4cf8e9a92a639845e0
- Sigstore transparency entry: 2256564578
- Sigstore integration time: Jul 26, 2026
Source repository:
- Permalink: nardi/splineax@eb54f866770ac42e1d6a700b34d665a0c4865259
- Branch / Tag: refs/tags/v0.2.2
- Owner: https://github.com/nardi
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: release.yml@eb54f866770ac42e1d6a700b34d665a0c4865259
- Trigger Event: release