Multi-Backend Support

Aurora-GLM supports three computational backends: NumPy (default), PyTorch, and JAX. All numerical code works transparently across backends.

Backend overview

Backend	Install	Strengths	When to use
NumPy	Core (default)	Stable, no extra deps	Default, small to medium data
PyTorch	pip install aurora-glm[torch]	GPU (CUDA), up to 141x speedup	Large datasets, GPU available
JAX	pip install aurora-glm[jax]	JIT compilation, automatic differentiation	Research, custom gradients

Automatic backend detection

The backend is detected automatically from input types:

import numpy as np
import torch
from aurora.models.glm import fit_glm

# NumPy arrays -> NumPy backend (automatic)
X_np = np.random.randn(1000, 5)
y_np = np.random.randn(1000)
result = fit_glm(X_np, y_np)

# PyTorch tensors -> PyTorch backend (automatic)
X_t = torch.tensor(X_np)
y_t = torch.tensor(y_np)
result = fit_glm(X_t, y_t)

The namespace pattern

Internal numerical code uses the multi-backend namespace pattern:

from aurora.distributions._utils import namespace, as_namespace_array

def my_function(x, y):
    # Detect backend from inputs
    xp = namespace(x, y)
    # Convert to appropriate array type
    x_arr = as_namespace_array(x, xp, like=y)
    # Use xp instead of np
    return xp.sum(x_arr * xp.exp(y))

This pattern ensures all operations go through xp (which is numpy, torch, or jax.numpy) instead of calling np.* directly.

GPU acceleration

PyTorch CUDA provides dramatic speedups for large problems:

import torch

# Move data to GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
X = torch.randn(100000, 50, device=device)
y = torch.randn(100000, device=device)

# Fit on GPU
from aurora.models.glm import fit_glm
result = fit_glm(X, y, family="gaussian")

When to use which backend

Scenario	Recommended	Reason
Default analysis	NumPy	No setup required, stable for n < 10,000
Large GLM (n > 10k)	PyTorch GPU	Matrix ops benefit from parallelism
GAM with many smooth terms	PyTorch GPU	Basis matrix construction is parallelizable
GAMM with many groups	PyTorch GPU	Random effects design matrix benefits from GPU
Custom likelihood / gradient	JAX	JIT + autodiff for custom objectives

Compatibility notes

• All backends produce numerically identical results (validated against R)

• PyTorch tensors are converted to NumPy internally where needed (e.g., for SciPy operations in GAMM)

• JAX arrays work similarly, with automatic JIT compilation where applicable

• Sparse matrix support is available on all backends via SciPy