kerneljax 0.1.0

A low-level JAX interface for nonparametric kernel smoothing of mixed-type data.

Docs · Quickstart · Background · API Reference

KernelJax is a low-level JAX library for nonparametric kernel smoothing of mixed-type data, built for researchers who develop new methodology from composable, differentiable building blocks. Its kernels, bandwidth selectors, and smoothers run natively on CPUs, GPUs, and TPUs through JAX and XLA, and fit naturally into the wider JAX ecosystem.

Features

• Estimators for samples mixing continuous, unordered categorical, and ordered categorical columns.
  • Density and cumulative distribution estimation.
  • Conditional density, distribution, quantiles, and modes given mixed covariates.
  • Local polynomial regression with derivatives and pointwise standard errors.
• Bandwidths and categorical smoothing parameters selected jointly from the data by least squares or likelihood cross validation, a corrected AIC, or a closed-form plug-in rule.
• Cross-validation criteria are ordinary JAX functions that compose with jax.grad, jax.jit, and jax.vmap, so a bandwidth can be learned inside a larger model rather than fixed by a derivative-free search.
• Composable primitives (kweights, ksum, the kernel base classes) are public, so estimators the high-level interface does not ship can be built directly.

Installation

uv pip install kerneljax

JAX itself is not pinned, so install the build matching your hardware first. The JAX installation guide covers the CPU, CUDA, and TPU wheels.

Quickstart

import numpy as np
import kerneljax as kj

rng = np.random.default_rng(1)
x = rng.uniform(size=200)
y = np.sin(2 * np.pi * x) + rng.normal(0, 0.2, 200)
fit = kj.local_poly(x, y, "cv_ls", degree=1)
print(kj.summary(fit))

Local polynomial regression

  Observations                         200
  Continuous variables                   1
  Estimator                   local linear
  Bandwidth type                    shared

  Variable      Kind             Bandwidth
  x1            continuous        0.036019

  Continuous kernel         Gaussian, order 2

  Residual standard error         0.174543
  R-squared                       0.947953

  Selection                          cv_ls
  Criterion value                 0.034301
  Solver iterations                      8
  Converged                           True

Custom kernels

A kernel sets how observations are weighted. Subclass ContinuousKernel or the categorical base for that column kind, implement value, and pass the result through KernelSet. Bandwidth selection then runs through the kernel you wrote.

import dataclasses
import jax.numpy as jnp

@dataclasses.dataclass(frozen=True)
class Epanechnikov(kj.ContinuousKernel):
    def value(self, x, y, h):
        u = (x - y) / h
        return jnp.where(jnp.abs(u) <= 1.0, 0.75 * (1.0 - u * u), 0.0)

epan = kj.local_poly(x, y, "cv_ls", degree=1, kernels=kj.KernelSet(continuous=Epanechnikov()))
print(f"Epanechnikov  h={epan.bandwidth.h[0]:.6f}  r2={epan.r_squared:.6f}")
print(f"Gaussian      h={fit.bandwidth.h[0]:.6f}  r2={fit.r_squared:.6f}")

Epanechnikov  h=0.084645  r2=0.947231
Gaussian      h=0.036019  r2=0.947953

Custom bandwidth selection

Bandwidth selection minimizes a scalar criterion. The built-in rules (cv_ls, corrected AIC) are ordinary JAX callables, and so is anything you write in their place. Implement __call__, pass it to select_bandwidth, and the optimizer finds h by differentiating the loss you wrote.

@dataclasses.dataclass(frozen=True)
class AbsoluteDeviation:
    degree: int = 1

    def __call__(self, train, bandwidth, *, y, kernels=None, chunk=None):
        fit = kj.local_poly(train, y, bandwidth, kernels=kernels, chunk=chunk,
                            degree=self.degree, fold=jnp.arange(train.n))
        return jnp.mean(jnp.abs(y - fit.mean))

squared = kj.select_bandwidth(x, kj.RegressionCriterion(method="cv_ls", degree=1), y=y)
absolute = kj.select_bandwidth(x, AbsoluteDeviation(degree=1), y=y)
print(f"squared error       h = {squared.bandwidth.h[0]:.4f}")
print(f"absolute deviation  h = {absolute.bandwidth.h[0]:.4f}")

squared error       h = 0.0360
absolute deviation  h = 0.0348