Metadata-Version: 2.4
Name: tnkm
Version: 0.1.0
Summary: Python library for tensor-network kernel machines powered by JAX.
License-Expression: MIT
License-File: LICENSE
Keywords: tensor-networks,kernel-methods,machine-learning,system-identification,tnkm,jax
Author: Albert Saiapin
Author-email: analystalb@gmail.com
Requires-Python: >=3.10
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Information Technology
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Dist: jax (>=0.6,<1)
Requires-Dist: numpy (>=2.2,<3)
Requires-Dist: optax (>=0.2.7,<1)
Project-URL: Documentation, https://albmlpy.github.io/tnkm/
Project-URL: Homepage, https://github.com/AlbMLpy/tnkm/
Description-Content-Type: text/markdown

<p align="center">
  <img
    src="https://raw.githubusercontent.com/AlbMLpy/tnkm/main/docs/assets/logo.png"
    width="160"
    alt="TNKM logo">
</p>

<h1 align='center'>Tensor Network Kernel Machines (TNKM)</h1>

[![Documentation](https://github.com/albmlpy/tnkm/actions/workflows/docs.yml/badge.svg)](https://albmlpy.github.io/tnkm/)
[![Test](https://github.com/albmlpy/tnkm/actions/workflows/tests.yml/badge.svg)](https://github.com/albmlpy/tnkm/actions/workflows/tests.yml)

TNKM is a [JAX](https://github.com/google/jax)-based library for machine learning with tensor-network parameterizations.

It combines kernel methods with low-rank tensor decompositions to build scalable models with explicit rank control.

TNKM is intended for research in machine learning, system identification, and time-series modeling.

## Features

- **CP and TT tensor-network kernel machines**
- **Polynomial**, **Fourier**, **B-spline** and other feature maps
- **Alternating Least Squares (ALS)** and **gradient-based optimization** ([Optax](https://github.com/google-deepmind/optax))
- **Explicit rank control** for balancing model complexity and expressiveness
- **JAX-native implementation** with JIT compilation and hardware acceleration

## Installation

### From PyPI

```bash
pip install tnkm
```

### From source

```bash
git clone https://github.com/AlbMLpy/tnkm.git
cd tnkm
pip install .
```

## Quick example

Fit a CP-based tensor-network kernel machine with polynomial features:

```python
import jax.numpy as jnp

from tnkm.models import CPKM
from tnkm.optim import train_als
from tnkm.features import ProductFeatures, PolyFeature

d_dim = 2
x = jnp.ones((32, d_dim))
y = jnp.ones((32,))

features = ProductFeatures((PolyFeature(4, k_col=i) for i in range(d_dim)))
model = CPKM(features, rank=2, seed=0)

train_als(model, x, y, n_epoch=3, gamma_w=1e-3, beta_e=1.0)
y_pred = model.predict(x)
```

## Documentation

Available at [https://albmlpy.github.io/tnkm](https://albmlpy.github.io/tnkm).

## Contributing

Contributions to `tnkm` are welcome. If you would like to report issues,
suggest improvements, or contribute code, please see
[CONTRIBUTING.md](https://github.com/AlbMLpy/tnkm/blob/main/CONTRIBUTING.md) for development guidelines,
installation instructions, and information about the contribution workflow.

## Citation

TNKM is an open-source research software project.
If you use TNKM in academic work, please cite the specific software version used.
Citation information will be updated when a DOI or publication becomes available.

(Also consider starring the project [on GitHub](https://github.com/AlbMLpy/tnkm))

## License

MIT License

