klujax

library0.5.0pypypiunverified

A JAX-compatible KLU sparse direct solver for linear systems (LU factorization), useful for circuit simulation and power systems. Version 0.5.0; pre-1.0, breaking changes possible. Cadence: irregular.

data ai-ml

pip install klujax

package ↗

INSTALL

IMPORT

SIG · KLUJAX

klujax

datapythonv0.5.0

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

klujax_solve

✓ from klujax import klujax_solve

✗ from klujax import klujax_solve

Solve a sparse linear system using KLU via JAX. Inputs are COO format arrays (row, col, data) and right-hand side b.

import jax.numpy as jnp from klujax import klujax_solve # Define a sparse matrix in COO format row = jnp.array([0, 1, 0, 1], dtype=jnp.int32) col = jnp.array([0, 0, 1, 1], dtype=jnp.int32) data = jnp.array([2.0, 1.0, 1.0, 3.0], dtype=jnp.float64) b = jnp.array([1.0, 2.0], dtype=jnp.float64) # Solve A x = b x = klujax_solve(row, col, data, b) print(x)

Debug

Known issues

breakingPre-1.0 API may change without notice. The function signature or return type may change in future releases.

fix

Pin version with `klujax==0.5.0` and monitor releases.

affects: all < 1.0

gotchaklujax_solve expects row, col, data as JAX int32/int64 and float32/float64 arrays. Using Python lists or incorrect dtypes may cause silent errors or type promotion issues.

fix

Ensure inputs are JAX arrays with correct dtypes: `row=jnp.array(..., dtype=jnp.int32)`, `data=jnp.array(..., dtype=jnp.float64)`.

affects: all

gotchaThe solver is not JIT-compatible by default; current implementation uses Python loops. JIT compilation may fail or produce incorrect results.

fix

Do not wrap klujax_solve inside `jax.jit`. Use outside JIT context.

affects: <=0.5.0

Upgrade

Version history

0.5.0latest on PyPI · released Apr 15, 2026

Audit

Dependencies

jaxrequiredCore dependency for JAX arrays and operations

jaxlibrequiredRequired for JAX runtime

Agent activity

9 hits · last 30 days

node

Amazon

Resources

packagehttps://pypi.org/project/klujax/ ↗