kernax-ml

KO | EN

gitlite — search

#deep-learning #markdown #chatgpt #css #awesome-list #cpp #c #privacy #openai #framework

(★ 9)

Kernax: a Python package providing efficient mathematical kernel implementations for probabilistic machine learning models, built with the JAX framework.

File Explorer

ZIP Download

- - publish.yml
  - tests.yml
- - benchmark_se_kernel.py
- __init__.py
- conftest.py
- input_generators.py
- README.md
- - computation_engines.md
  - formatting.md
  - reparameterization.md
- - index.md
- - img.png
- - changelog.md
  - contributing_guide.md
- - gp_regression.py
  - kernel_composition.py
  - kernel_ridge_regression.py
  - multioutput_convolution.py
  - multioutput_icm.py
  - multioutput_lcm.py
  - svm.py
  - svr.py
- - all_kernels.md
  - installation.md
  - key_concepts.md
  - quick_start.md
  - sharp_bits.md
  - should_i_use_kernax.md
- - batching_modules.md
  - benchmarking.md
  - custom_modules.md
  - freezing_parameters.md
  - mutate_parameters.md
  - saving_modules.md
  - using_priors.md
- .gitignore
- about.md
- citation.md
- conf.py
- faq.md
- index.md
- - __init__.py
  - AffineKernel.py
  - DotProductKernel.py
  - LinearKernel.py
  - PolynomialKernel.py
  - Sigmoid.py
- - __init__.py
  - AffineMean.py
  - ConstantMean.py
  - LinearMean.py
  - ZeroMean.py
- - __init__.py
  - _gather.py
  - BlockDiagKernel.py
  - BlockMean.py
  - ConvolutionKernel.py
  - ICMKernel.py
  - LCMKernel.py
- - __init__.py
  - AbstractOperatorModule.py
  - ProductModule.py
  - SumModule.py
- - __init__.py
  - ConstantKernel.py
  - VarianceKernel.py
- - __init__.py
  - Matern12Kernel.py
  - Matern32Kernel.py
  - Matern52Kernel.py
  - PeriodicKernel.py
  - RationalQuadraticKernel.py
  - SEKernel.py
  - StationaryKernel.py
  - WhiteNoiseKernel.py
- - __init__.py
  - ActiveDimsModule.py
  - ARDKernel.py
  - BatchModule.py
  - ExpModule.py
  - InputSpecificParamModule.py
  - LogModule.py
  - NegModule.py
  - WrapperModule.py
- __init__.py
- AbstractKernel.py
- AbstractMean.py
- distances.py
- engines.py
- hp_sampling.py
- mask.py
- module.py
- parametrisations.py
- py.typed
- types.py
- utils.py
- README.md
- setup_dev.sh
- verify_setup.py
- __init__.py
- allurerc.json
- block.ipynb
- conftest.py
- README.md
- test_hp_sampling.py
- test_kernel_basics.py
- test_kernel_compositions.py
- test_kernel_formatting.ipynb
- test_kernel_mutations.py
- test_kernel_wrappers.py
- test_mean_basics.py
- test_mean_compositions.py
- test_mean_formatting.ipynb
- test_mean_mutations.py
- test_mean_wrappers.py
- test_multioutput.py
- test_parametrisations.py
.editorconfig
.gitignore
.readthedocs.yaml
CHANGELOG.md
CLAUDE.md
CONTRIBUTING.md
DOC_STRUCTURE.MD
environment.yml
LICENSE
Makefile
MANIFEST.in
pyproject.toml
README.md

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

Kernax is a Python package providing **efficient mathematical kernel implementations** for *probabilistic machine 
learning models*, built with the **JAX framework**.

Kernels are critical elements of probabilistic models, used in many inner-most loops to compute giant matrices and 
optimise numerous hyper-parameters. Therefore, this library emphasise on efficient, modular and scalable kernel 
implementations, with the following features:

- **JIT-compiled** computations for fast execution on CPU, GPU and TPU
- Kernels structured as **Equinox Modules** (aka *PyTrees*) which means...
  - They can be sent to (jitted) functions **as parameters**
  - Their **hyper-parameters can be optimised via autodiff**
  - They can **be vectorised-on** with `vmap`
- **Composable kernels** through operator overloading (`+`, `*`, `-`)
- **Kernel wrappers** to scale to higher dimensions (batch or block of covariance matrices)
- **Multi-output kernels** (ICM, LCM, convolution processes) for correlated GP outputs
- **NaN-aware computations** for working with padded/masked data

> **⚠️ Project Status**: Kernax is in early development. The API may change, and some features are still experimental.

## Installation

Install from PyPI:

```bash
pip install kernax-ml
```

Or clone the repository for development:

```bash
git clone https://github.com/SimLej18/kernax-ml
cd kernax-ml
```

**Requirements**:
- Python >= 3.10
- JAX >= 0.6.2

**Using Conda** (recommended):

```bash
conda create -n kernax-ml python=3.12
conda activate kernax-ml
pip install -e .
```

**Using pip**:

```bash
pip install -e .
```

## Quick Start

```python
import jax.numpy as jnp
from kernax import SEKernel, LinearKernel, WhiteNoiseKernel, ExpModule, BatchModule, ARDKernel

# Create a simple Squared Exponential kernel
kernel = SEKernel(length_scale=1.0)

# Compute covariance between two points
x1 = jnp.array([1.0, 2.0])
x2 = jnp.array([1.5, 2.5])
cov = kernel(x1, x2)

# Compute covariance matrix for a set of points
X = jnp.array([[1.0], [2.0], [3.0]])
K = kernel(X, X)  # Returns 3x3 covariance matrix

# Compose kernels using operators
composite_kernel = SEKernel(length_scale=1.0) + WhiteNoiseKernel(0.1)  # SE + noise

# Use BatchKernel for distinct hyperparameters per batch
base_kernel = SEKernel(length_scale=1.0)
batched_kernel = BatchModule(base_kernel, batch_size=10, batch_in_axes=0, batch_over_inputs=True)

# Use ARDKernel for Automatic Relevance Determination
length_scales = jnp.array([1.0, 2.0, 0.5])  # Different scale per dimension
ard_kernel = ARDKernel(SEKernel(length_scale=1.0), length_scales=length_scales)
```

## Available Kernels

### Base Kernels

- **`SEKernel`** (Squared Exponential, aka RBF or Gaussian)
  - Hyperparameters: `length_scale`

- **`LinearKernel`**
  - Hyperparameters: `variance_b`, `variance_v`, `offset_c`

- **`MaternKernel`** family
  - `Matern12Kernel` (ν=1/2, equivalent to Exponential)
  - `Matern32Kernel` (ν=3/2)
  - `Matern52Kernel` (ν=5/2)
  - Hyperparameters: `length_scale`

- **`PeriodicKernel`**
  - Hyperparameters: `length_scale`, `variance`, `period`

- **`RationalQuadraticKernel`**
  - Hyperparameters: `length_scale`, `variance`, `alpha`

- **`ConstantKernel`**
  - Hyperparameters: `value`

- **`PolynomialKernel`**
  - Hyperparameters: `degree`, `gamma`, `constant`

- **`SigmoidKernel`** (Hyperbolic Tangent)
  - Hyperparameters: `alpha`, `constant`

- **`WhiteNoiseKernel`**
  - Diagonal noise kernel (returns constant value only on diagonal)
  - Implemented as `ConstantKernel` with `SafeDiagonalEngine`
  - Hyperparameters: `noise`

### Composite Modules

- **`SumModule`**: Adds two modules (use `kernel1 + kernel2`)
- **`ProductModule`**: Multiplies two modules (use `kernel1 * kernel2`)

### Wrapper Modules

Transform or modify kernel behavior:

- **`ExpModule`**: Applies exponential to kernel output
- **`LogModule`**: Applies logarithm to kernel output
- **`NegModule`**: Negates kernel output (use `-kernel`)
- **`BatchModule`**: Adds batch handling with distinct hyperparameters per batch
- **`ActiveDimsModule`**: Selects specific input dimensions before kernel computation
- **`ARDKernel`**: Applies Automatic Relevance Determination (different length scale per dimension)

### Multi-output Kernels

For correlated outputs in multi-output GPs (`kernax.multioutput`). `output_ids` (+ `output_ids2`)
is passed to `__call__`, not stored on the kernel — supports both a grid shared across outputs
and heterotopic data (arbitrary per-output sizes/order):

- **`ICMKernel`**: Intrinsic Coregionalisation Model, `K = B ⊗ k(x1, x2)` with `B = W Wᵀ + diag(κ)`
- **`LCMKernel`**: Linear Coregionalisation Model, a sum of independent `ICMKernel`s
- **`ConvolutionKernel`**: convolution-process kernel with per-output bandwidth/variance
- **`BlockDiagKernel`** / **`BlockMean`**: independent per-output kernel/mean, no cross-output correlation

### Computation Engines

Computation engines control how covariance matrices are computed. All kernels accept an `engine` parameter:

- **`DenseEngine`** (default): Computes full covariance matrices
- **`NaNDenseEngine`**: applies a shortcut to skip computing NaN vectors
- **`MaskedNaNEngine`**: masks NaN entries directly in the kernel matrix computation

Example:
```python
from kernax import SEKernel
from kernax.engines import NaNDenseEngine

# Skip NaN vectors when computing the covariance matrix
diagonal_kernel = SEKernel(length_scale=1.0, engine=NaNDenseEngine)
```

## Architecture

Kernax is built on [Equinox](https://github.com/patrick-kidger/equinox), so they are compatible with every feature from JAX!

Each kernel is a single class extending `AbstractKernel` (which extends `eqx.Module`). It implements `pairwise(self, x1, x2)` for scalar covariance computation, and delegates matrix construction to an interchangeable `engine`. Hyperparameters are stored in wrapped form and exposed via properties, with a per-hyperparameter `AbstractParametrisation` controlling their transformation during optimization.

## Testing & Quality

Kernax maintains high code quality standards:

- **338 passing tests**
- **Allure test reporting** for detailed test analytics
- **Cross-library validation** against scikit-learn, GPyTorch, and GPJax
- **Type checking** with mypy for enhanced code safety
- **Code formatting** with ruff (tabs, line length 100)

Run tests with:
```bash
make test           # Run all tests
make test-cov       # Run tests with coverage report
make test-allure    # Generate Allure HTML report
make lint           # Run type checking and linting
```

## Benchmarks

Kernax is designed for performance. You can run a benchmark comparison with other libraries with:
```bash
make benchmarks-compare
```

Our preliminary results show a significant speed-up over alternatives when JIT compilation is enabled:
```
------------------------------------------------- benchmark 'benchmarks/comparison/compare_se_kernel.py::Benchmark1DRandom::test_compare': 4 tests ------------------------------------------------
Name (time in ms)               Min                 Max                Mean            StdDev              Median               IQR            Outliers  OPS (mops/s)            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_compare[kernax]        12.0889 (1.0)       14.9765 (1.0)       12.6690 (1.0)      0.6474 (1.0)       12.4472 (1.0)      0.5749 (1.29)          1;1   78,932.9054 (1.0)          20           1
test_compare[gpytorch]      43.5577 (3.60)      54.3097 (3.63)      44.5481 (3.52)     2.3159 (3.58)      43.9814 (3.53)     0.4448 (1.0)           1;1   22,447.6602 (0.28)         20           1
test_compare[gpjax]         67.3657 (5.57)      73.9067 (4.93)      68.8340 (5.43)     1.4448 (2.23)      68.5019 (5.50)     1.3212 (2.97)          3;1   14,527.6964 (0.18)         20           1
test_compare[sklearn]      328.1409 (27.14)    367.2989 (24.53)    334.7924 (26.43)    8.2573 (12.75)    332.5784 (26.72)    4.2589 (9.57)          1;1    2,986.9256 (0.04)         20           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

---------------------------------------------- benchmark 'benchmarks/comparison/compare_se_kernel.py::Benchmark1DRegularGrid::test_compare': 4 tests -----------------------------------------------
Name (time in ms)               Min                 Max                Mean             StdDev              Median               IQR            Outliers  OPS (mops/s)            Rounds  Iterations
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_compare[kernax]        11.8415 (1.0)       13.3571 (1.0)       12.5704 (1.0)       0.4358 (1.0)       12.5185 (1.0)      0.6904 (2.23)          8;0   79,551.8567 (1.0)          20           1
test_compare[gpytorch]      43.6603 (3.69)      55.0724 (4.12)      44.5337 (3.54)      2.4908 (5.72)      43.9668 (3.51)     0.3099 (1.0)           1;1   22,454.9091 (0.28)         20           1
test_compare[gpjax]         67.2976 (5.68)     119.1254 (8.92)      70.9630 (5.65)     11.3640 (26.08)     68.3379 (5.46)     0.8114 (2.62)          1;2   14,091.8552 (0.18)         20           1
test_compare[sklearn]      297.8652 (25.15)    316.3752 (23.69)    302.3710 (24.05)     4.4811 (10.28)    300.7931 (24.03)    3.5605 (11.49)         4;2    3,307.1952 (0.04)         20           1
----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

---------------------------------------------- benchmark 'benchmarks/comparison/compare_se_kernel.py::Benchmark2DMissingValues::test_compare': 4 tests ----------------------------------------------
Name (time in ms)               Min                 Max                Mean             StdDev              Median                IQR            Outliers  OPS (mops/s)            Rounds  Iterations
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_compare[kernax]        11.9619 (1.0)       13.9954 (1.0)       12.7085 (1.0)       0.5785 (1.0)       12.4387 (1.0)       0.8605 (1.0)           5;0   78,687.2477 (1.0)          20           1
test_compare[gpytorch]      25.9657 (2.17)      30.2475 (2.16)      27.1899 (2.14)      1.3834 (2.39)      26.6297 (2.14)      1.4433 (1.68)          4;2   36,778.4048 (0.47)         20           1
test_compare[gpjax]         55.1528 (4.61)     136.3582 (9.74)     113.5941 (8.94)     26.1170 (45.15)    125.4449 (10.09)    16.2559 (18.89)         3;3    8,803.2728 (0.11)         20           1
test_compare[sklearn]      213.0630 (17.81)    272.6552 (19.48)    230.3581 (18.13)    15.2409 (26.35)    224.1107 (18.02)    14.8310 (17.23)         6;1    4,341.0674 (0.06)         20           1
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

------------------------------------------------- benchmark 'benchmarks/comparison/compare_se_kernel.py::Benchmark2DRandom::test_compare': 4 tests ------------------------------------------------
Name (time in ms)               Min                 Max                Mean            StdDev              Median               IQR            Outliers  OPS (mops/s)            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_compare[kernax]        13.9697 (1.0)       15.5467 (1.0)       14.5331 (1.0)      0.4723 (1.0)       14.3701 (1.0)      0.8122 (1.57)          8;0   68,808.3748 (1.0)          20           1
test_compare[gpytorch]      43.6454 (3.12)      50.5380 (3.25)      44.5098 (3.06)     1.4877 (3.15)      44.1243 (3.07)     0.5169 (1.0)           1;2   22,466.9466 (0.33)         20           1
test_compare[gpjax]         94.1932 (6.74)     103.3563 (6.65)      97.6833 (6.72)     2.1244 (4.50)      97.8704 (6.81)     3.0956 (5.99)          4;0   10,237.1672 (0.15)         20           1
test_compare[sklearn]      408.6985 (29.26)    440.7834 (28.35)    417.1601 (28.70)    8.0450 (17.03)    413.5329 (28.78)    9.0707 (17.55)         5;1    2,397.1614 (0.03)         20           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

---------------------------------------------- benchmark 'benchmarks/comparison/compare_se_kernel.py::Benchmark2DRegularGrid::test_compare': 4 tests ----------------------------------------------
Name (time in ms)               Min                 Max                Mean            StdDev              Median               IQR            Outliers  OPS (mops/s)            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_compare[kernax]        13.9134 (1.0)       16.3335 (1.0)       14.5133 (1.0)      0.7465 (1.0)       14.1642 (1.0)      0.8484 (1.0)           3;2   68,902.5346 (1.0)          20           1
test_compare[gpytorch]      43.9418 (3.16)      51.8444 (3.17)      45.7195 (3.15)     1.8991 (2.54)      45.4059 (3.21)     1.9106 (2.25)          2;2   21,872.5192 (0.32)         20           1
test_compare[gpjax]         93.1488 (6.69)     130.6572 (8.00)     102.3413 (7.05)     8.1851 (10.96)    101.9908 (7.20)     9.7936 (11.54)         2;1    9,771.2271 (0.14)         20           1
test_compare[sklearn]      381.4240 (27.41)    405.2938 (24.81)    387.4473 (26.70)    6.4979 (8.70)     384.8169 (27.17)    9.0655 (10.69)         2;0    2,580.9964 (0.04)         20           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
```

## Development Status

Check the [changelog](CHANGELOG.md) for details.

### ✅ Completed

- Core kernel implementations (SE, Linear, Matern, Periodic, Sigmoid, etc.)
- Kernel composition via operators
- Equinox Module integration
- NaN-aware computations
- BatchKernel wrapper with distinct/shared hyper-parameters
- ARDKernel wrapper using input scaling
- ActiveDimsKernel wrapper for dimension selection
- Multi-output kernels (ICM, LCM, convolution process) with `BlockDiagKernel`/`BlockMean`
- StationaryKernel and DotProductKernel base classes with proper inheritance
- Per-hyperparameter parametrisation system (LogExp, Softplus, Bounded, NonTrainable)
- Parameter freezing either via `NonTrainableParametrisation` or via masking and using eqx.partition+eqx.combine
- Computation engines for special cases (NaN-aware)
- Comprehensive test suite
- Benchmark architecture
- PyPI package distribution

### 🚧 In Progress / Planned

- Expanded documentation and tutorials

## Contributing

This project is in early development. Contributions, bug reports, and feature requests are welcome!

## Related Projects

Kernax is developed alongside [MagmaClust](https://github.com/SimLej18/MagmaClustPy), a clustering and Gaussian Process library.

## License

MIT License - see [LICENSE](LICENSE) file for details.

## Citation

[Citation information to be added]

kernax-ml

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

kernax-ml

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

Download README