cholops

Cholesky modifications in JAX, with differentiable tensor operations and immutable factor objects. cholops provides fast appends, deletions and updates of a dense positive-definite matrix's factor, directly or through a factor object.

Install

Install with pip install cholops. On Linux with a compatible CUDA 13 driver, use pip install "cholops[cuda13]". The initial release targets JAX 0.11.2 and Python 3.12 or newer.

We recommend using float64; see JAX's precision settings.

Set the JAX_ENABLE_X64 shell environment variable before starting Python:

export JAX_ENABLE_X64=1

Alternatively, enable it in Python before creating arrays:

import jax

jax.config.update("jax_enable_x64", True)

Example

import jax
import jax.numpy as jnp
from cholops import Cholesky, PermutedCholesky

A = jnp.array([[2.0, 0.3, 0.1], [0.3, 1.5, 0.2], [0.1, 0.2, 1.2]])
b = jnp.array([1.0, -0.5, 0.2])

# Create a Cholesky factor from a positive-definite matrix.
factor = Cholesky.from_matrix(A)
factor = factor.append(jnp.array([0.2, 0.1, 0.05]), 1.3)
L = factor.factor
x = factor.solve(
    jnp.append(b, 0.4)
)  # One RHS entry per variable, including the new one.

# PermutedCholesky replaces by deletion and appending,
# while managing logical ordering transparently.
managed = PermutedCholesky.from_matrix(A)
updated = managed.replace(1, jnp.array([0.2, 0.1]), 1.3)
x_new = updated.solve(b)  # Inputs and outputs remain in logical order.


# Zero-mean Gaussian log likelihood and its covariance/data gradients.
@jax.jit
def log_likelihood(A, b):
    state = Cholesky.from_matrix(A)
    return -0.5 * (state.logdet() + state.quadratic(b) + b.size * jnp.log(2 * jnp.pi))


logp, (grad_A, grad_b) = jax.value_and_grad(log_likelihood, argnums=(0, 1))(A, b)

See the factor guide for more information. The executable examples explain each step and check their results against dense calculations:

• Basic factor operations: add, replace and remove observations, then solve a system and calculate a Gaussian log likelihood.
• Streaming Gaussian process: maintain a fixed-size observation window with PermutedCholesky and differentiate through the replacement sequence with respect to a kernel parameter.
• Online least squares: update a regression fit with rank1_positive as measurements arrive.
• Direct tensor operations: compare replacement orderings and arrange the right-hand side for each returned factor.

Tensor operations

The factor objects use the following tensor operations, which can also be imported directly from cholops:

See the tensor guide for shapes, ordering and preconditions. All operations support jit, vmap and first-order forward and reverse differentiation. Higher-order derivatives have not been validated for general use.

Performance

We benchmark a single index-preserving chol_replace call against a naive JAX implementation and full matrix refactorization on an NVIDIA H100. All methods receive the existing factor and replacement covariances.

Replacing the first index of a 16,384 × 16,384 matrix gives the following speedups after compilation:

Recomputation forms L @ L.T, changes its row/column and calls native JAX Cholesky. The naive replacement baseline uses two signed rank-one scans with standard JAX autodiff.

See the complete benchmark results for the measurement protocol, additional timings and memory usage.

Numerical limitations

Near-singular covariances can produce inaccurate factors and gradients. See numerical limits.