mlx-fisher 0.1.0

GPU-accelerated Fisher Information Matrix computation on Apple Silicon via MLX

mlx-fisher

GPU-accelerated Fisher Information Matrix computation on Apple Silicon via MLX.

Author: Sheng-Kai Huang ([e-mail])

Features

• Fisher Information Matrix from log-likelihood, model predictions, or samples
• CMB C_l Fisher matrix for cosmological parameter estimation
• KL divergence (classical) and quantum relative entropy D(rho||sigma)
• Natural gradient descent optimizer with online Fisher estimation
• All matrix operations (eigendecomposition, matrix multiply, log) on Apple GPU

Installation

pip install -e .

Requires Python 3.10+ and Apple Silicon (M1/M2/M3/M4).

Quick Start

import mlx.core as mx
from mlx_fisher import FisherMatrix, kl_divergence, quantum_relative_entropy

# --- Gaussian Fisher matrix from a model ---
x = mx.linspace(-3.0, 3.0, 1000)

def model(theta):
    return theta[0] * x**2 + theta[1] * x + theta[2]

theta0 = mx.array([1.0, -0.5, 2.0])
sigma = mx.ones((1000,)) * 0.5

F = FisherMatrix.from_model(model, theta0, sigma)
print(F.marginal_errors())   # 1-sigma errors on each parameter

# --- KL divergence ---
p = mx.array([0.4, 0.3, 0.2, 0.1])
q = mx.array([0.25, 0.25, 0.25, 0.25])
print(kl_divergence(p, q))   # D_KL(p || q)

# --- Quantum relative entropy ---
d = 4
rho = mx.zeros((d, d)); rho[0, 0] = 1.0       # pure state |0><0|
sigma_dm = mx.eye(d) / d                        # maximally mixed
print(quantum_relative_entropy(rho, sigma_dm))   # = ln(d)

CMB Fisher Matrix

from mlx_fisher import fisher_matrix_cl

def cl_fn(theta):
    """Map cosmological parameters to C_l power spectrum."""
    # Your Boltzmann solver here (e.g., CLASS wrapper)
    ...

theta_fid = mx.array([0.022, 0.12, 0.06, 0.96, 3.04, 67.4])
F = fisher_matrix_cl(cl_fn, theta_fid, f_sky=0.7, l_min=2, l_max=2500)
print(F.marginal_errors())

Natural Gradient Descent

from mlx_fisher import NaturalGradientOptimizer

opt = NaturalGradientOptimizer(lr=1e-2, damping=1e-4)

for step in range(100):
    grad = compute_gradient(theta)
    theta = opt.step(theta, grad, fisher_estimator=compute_fisher)

API Reference

FisherMatrix

• FisherMatrix.from_model(model_fn, theta, sigma) -- Gaussian Fisher matrix
• FisherMatrix.from_loglikelihood(log_lik, theta) -- from log-likelihood function
• FisherMatrix.from_samples(log_prob, theta, samples) -- empirical Fisher
• .inverse(reg=0.0) -- covariance matrix (regularised inversion)
• .marginal_errors(reg=0.0) -- 1-sigma marginal errors
• .eigenvalues() -- eigenvalue spectrum
• .condition_number() -- matrix condition number

fisher_matrix_cl(cl_fn, theta, f_sky, noise_cl, l_min, l_max)

CMB power spectrum Fisher matrix with cosmic variance.

kl_divergence(p, q)

Classical KL divergence D_KL(p || q).

quantum_relative_entropy(rho, sigma)

Quantum relative entropy D(rho || sigma) = Tr[rho(ln rho - ln sigma)].

NaturalGradientOptimizer(lr, damping, fisher_update_interval, ema_decay)

Natural gradient descent: theta_new = theta - lr * F^{-1} @ grad.

Benchmarks (M1 Max)

Operation	Scale	MLX (ms)	NumPy (ms)	Speedup
KL divergence	1M bins	0.52	5.22	10x
KL divergence	10M bins	2.23	51.40	23x
Eigendecomposition	512x512	17.25	58.07	3.4x
Matrix multiply (Fisher)	32768x512	4.33	169.96	39x

See benchmark_results.md for full results.

License

MIT