qjax 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- qjax-0.1.0/.coverage +0 -0
- qjax-0.1.0/.github/workflows/ci.yaml +56 -0
- qjax-0.1.0/.github/workflows/release.yaml +27 -0
- qjax-0.1.0/.gitignore +15 -0
- qjax-0.1.0/.plan/PLAN.md +171 -0
- qjax-0.1.0/.readthedocs.yaml +16 -0
- qjax-0.1.0/CONTRIBUTING.md +131 -0
- qjax-0.1.0/LICENSE +11 -0
- qjax-0.1.0/PKG-INFO +195 -0
- qjax-0.1.0/README.md +174 -0
- qjax-0.1.0/assets/favicon.svg +32 -0
- qjax-0.1.0/assets/logo_qjax.svg +34 -0
- qjax-0.1.0/docs/Makefile +16 -0
- qjax-0.1.0/docs/_static/custom.css +11 -0
- qjax-0.1.0/docs/_static/examples/attention_mlp.png +0 -0
- qjax-0.1.0/docs/_static/examples/attention_q_learning.gif +0 -0
- qjax-0.1.0/docs/_static/examples/classification.png +0 -0
- qjax-0.1.0/docs/_static/examples/classification_boundaries.png +0 -0
- qjax-0.1.0/docs/_static/examples/learnable_q.png +0 -0
- qjax-0.1.0/docs/_static/examples/optimization.gif +0 -0
- qjax-0.1.0/docs/_static/examples/optimization.png +0 -0
- qjax-0.1.0/docs/_static/examples/q_gaussian.png +0 -0
- qjax-0.1.0/docs/_static/examples/reinforcement_learning.png +0 -0
- qjax-0.1.0/docs/_static/favicon.svg +32 -0
- qjax-0.1.0/docs/_static/logo.svg +34 -0
- qjax-0.1.0/docs/_templates/.gitkeep +0 -0
- qjax-0.1.0/docs/api.md +52 -0
- qjax-0.1.0/docs/conf.py +81 -0
- qjax-0.1.0/docs/examples/attention_mlp.md +64 -0
- qjax-0.1.0/docs/examples/attention_q_learning.md +52 -0
- qjax-0.1.0/docs/examples/classification.md +92 -0
- qjax-0.1.0/docs/examples/learnable_q.md +63 -0
- qjax-0.1.0/docs/examples/optimization.md +67 -0
- qjax-0.1.0/docs/examples/q_gaussian.md +68 -0
- qjax-0.1.0/docs/examples/reinforcement_learning.md +70 -0
- qjax-0.1.0/docs/examples.md +65 -0
- qjax-0.1.0/docs/index.md +144 -0
- qjax-0.1.0/docs/installation.md +64 -0
- qjax-0.1.0/docs/quickstart.md +67 -0
- qjax-0.1.0/docs/theory.md +104 -0
- qjax-0.1.0/examples/attention_mlp.py +281 -0
- qjax-0.1.0/examples/attention_q_learning.py +212 -0
- qjax-0.1.0/examples/classification.py +379 -0
- qjax-0.1.0/examples/figures/attention_mlp.pdf +0 -0
- qjax-0.1.0/examples/figures/attention_q_learning.gif +0 -0
- qjax-0.1.0/examples/figures/attention_q_learning.pdf +0 -0
- qjax-0.1.0/examples/figures/classification.pdf +0 -0
- qjax-0.1.0/examples/figures/classification_boundaries.pdf +0 -0
- qjax-0.1.0/examples/figures/learnable_q.pdf +0 -0
- qjax-0.1.0/examples/figures/optimization.gif +0 -0
- qjax-0.1.0/examples/figures/optimization.pdf +0 -0
- qjax-0.1.0/examples/figures/q_gaussian.pdf +0 -0
- qjax-0.1.0/examples/figures/reinforcement_learning.pdf +0 -0
- qjax-0.1.0/examples/learnable_q.py +69 -0
- qjax-0.1.0/examples/optimization.py +175 -0
- qjax-0.1.0/examples/q_gaussian.py +48 -0
- qjax-0.1.0/examples/reinforcement_learning.py +195 -0
- qjax-0.1.0/pyproject.toml +51 -0
- qjax-0.1.0/qjax/__init__.py +52 -0
- qjax-0.1.0/qjax/core/__init__.py +36 -0
- qjax-0.1.0/qjax/core/activations.py +111 -0
- qjax-0.1.0/qjax/core/distributions.py +155 -0
- qjax-0.1.0/qjax/core/entropy.py +108 -0
- qjax-0.1.0/qjax/core/functions.py +161 -0
- qjax-0.1.0/qjax/plots/__init__.py +15 -0
- qjax-0.1.0/qjax/plots/distributions.py +47 -0
- qjax-0.1.0/qjax/plots/functions.py +71 -0
- qjax-0.1.0/qjax/plots/style.py +116 -0
- qjax-0.1.0/qjax/shared/__init__.py +6 -0
- qjax-0.1.0/qjax/shared/types.py +20 -0
- qjax-0.1.0/qjax/shared/validation.py +45 -0
- qjax-0.1.0/tests/test_activations.py +69 -0
- qjax-0.1.0/tests/test_distributions.py +83 -0
- qjax-0.1.0/tests/test_entropy.py +111 -0
- qjax-0.1.0/tests/test_functions.py +71 -0
- qjax-0.1.0/uv.lock +2104 -0
qjax-0.1.0/.coverage
ADDED
|
Binary file
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [main]
|
|
8
|
+
workflow_call:
|
|
9
|
+
|
|
10
|
+
concurrency:
|
|
11
|
+
group: ${{ github.workflow }}-${{ github.ref }}
|
|
12
|
+
cancel-in-progress: true
|
|
13
|
+
|
|
14
|
+
jobs:
|
|
15
|
+
lint:
|
|
16
|
+
runs-on: ubuntu-latest
|
|
17
|
+
steps:
|
|
18
|
+
- uses: actions/checkout@v4
|
|
19
|
+
- name: Set up uv
|
|
20
|
+
uses: astral-sh/setup-uv@v4
|
|
21
|
+
with:
|
|
22
|
+
enable-cache: true
|
|
23
|
+
- name: Install dependencies
|
|
24
|
+
run: uv sync --extra dev --python 3.12
|
|
25
|
+
- name: Lint with ruff
|
|
26
|
+
run: uv run ruff check qjax tests examples
|
|
27
|
+
|
|
28
|
+
test:
|
|
29
|
+
runs-on: ubuntu-latest
|
|
30
|
+
strategy:
|
|
31
|
+
fail-fast: false
|
|
32
|
+
matrix:
|
|
33
|
+
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
|
34
|
+
steps:
|
|
35
|
+
- uses: actions/checkout@v4
|
|
36
|
+
- name: Set up uv
|
|
37
|
+
uses: astral-sh/setup-uv@v4
|
|
38
|
+
with:
|
|
39
|
+
enable-cache: true
|
|
40
|
+
- name: Install dependencies
|
|
41
|
+
run: uv sync --extra dev --python ${{ matrix.python-version }}
|
|
42
|
+
- name: Run tests
|
|
43
|
+
run: uv run pytest --cov=qjax --cov-report=term-missing
|
|
44
|
+
|
|
45
|
+
docs:
|
|
46
|
+
runs-on: ubuntu-latest
|
|
47
|
+
steps:
|
|
48
|
+
- uses: actions/checkout@v4
|
|
49
|
+
- name: Set up uv
|
|
50
|
+
uses: astral-sh/setup-uv@v4
|
|
51
|
+
with:
|
|
52
|
+
enable-cache: true
|
|
53
|
+
- name: Install dependencies
|
|
54
|
+
run: uv sync --extra docs --python 3.12
|
|
55
|
+
- name: Build documentation (warnings as errors)
|
|
56
|
+
run: uv run sphinx-build -b html docs docs/_build/html -W --keep-going
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
name: Release
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
workflow_dispatch:
|
|
5
|
+
inputs:
|
|
6
|
+
version:
|
|
7
|
+
description: "Version tag (e.g., 0.1.0)"
|
|
8
|
+
required: true
|
|
9
|
+
type: string
|
|
10
|
+
|
|
11
|
+
jobs:
|
|
12
|
+
ci:
|
|
13
|
+
uses: ./.github/workflows/ci.yaml
|
|
14
|
+
|
|
15
|
+
publish:
|
|
16
|
+
needs: ci
|
|
17
|
+
runs-on: ubuntu-latest
|
|
18
|
+
steps:
|
|
19
|
+
- uses: actions/checkout@v4
|
|
20
|
+
- name: Set up uv
|
|
21
|
+
uses: astral-sh/setup-uv@v4
|
|
22
|
+
- name: Build package
|
|
23
|
+
run: uv build
|
|
24
|
+
- name: Publish to PyPI
|
|
25
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
26
|
+
with:
|
|
27
|
+
password: ${{ secrets.PYPI_API_TOKEN }}
|
qjax-0.1.0/.gitignore
ADDED
qjax-0.1.0/.plan/PLAN.md
ADDED
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
# qjax — Implementation Plan
|
|
2
|
+
|
|
3
|
+
> **qjax** is a research-oriented Python library built on [JAX](https://github.com/google/jax)
|
|
4
|
+
> whose mission is to spread the research and capabilities of **Tsallis statistics**
|
|
5
|
+
> (non-extensive statistical mechanics) in **artificial intelligence**.
|
|
6
|
+
|
|
7
|
+
## 1. Vision
|
|
8
|
+
|
|
9
|
+
Tsallis statistics generalizes Boltzmann–Gibbs statistics through a single
|
|
10
|
+
*entropic index* `q`. As `q → 1` every construction collapses back to its classical
|
|
11
|
+
counterpart (Shannon entropy, Gaussian, softmax, KL divergence, …). `qjax` exposes
|
|
12
|
+
these `q`-deformed primitives as **pure, differentiable, `jit`/`vmap`-friendly JAX
|
|
13
|
+
functions** so that researchers can:
|
|
14
|
+
|
|
15
|
+
- treat `q` as a *learnable parameter* and let models discover their own statistics;
|
|
16
|
+
- explore heavy-tailed / sparse alternatives to standard AI building blocks;
|
|
17
|
+
- reproduce and extend Tsallis-statistics results across ML domains.
|
|
18
|
+
|
|
19
|
+
## 2. Design principles
|
|
20
|
+
|
|
21
|
+
1. **Purity & composability** — every primitive is a pure function of `(x, q)`;
|
|
22
|
+
no hidden state. Everything works under `jax.jit`, `jax.grad`, `jax.vmap`.
|
|
23
|
+
2. **Numerically safe `q → 1` limit** — use the double-`where` trick so gradients
|
|
24
|
+
are finite exactly at `q = 1`.
|
|
25
|
+
3. **Single source of truth** — `core/` holds the math; `plots/`, `examples/`,
|
|
26
|
+
and `tests/` only consume the public API.
|
|
27
|
+
4. **Precise docstrings** — every public function documents its math (with the
|
|
28
|
+
defining formula and the `q → 1` limit), args, shapes, and returns.
|
|
29
|
+
5. **Consistent visuals** — all plots use the **`magma`** colormap via a single
|
|
30
|
+
style module.
|
|
31
|
+
|
|
32
|
+
## 3. Target package structure
|
|
33
|
+
|
|
34
|
+
```
|
|
35
|
+
qjax/
|
|
36
|
+
├── __init__.py # curated public API
|
|
37
|
+
├── core/
|
|
38
|
+
│ ├── __init__.py # re-exports core primitives
|
|
39
|
+
│ ├── functions.py # q_log, q_exp + q-algebra (q_add, q_diff, q_prod, q_div)
|
|
40
|
+
│ ├── entropy.py # tsallis_entropy, tsallis_cross_entropy, tsallis_divergence
|
|
41
|
+
│ ├── distributions.py # q-Gaussian: pdf, logpdf, sample, normalization C_q
|
|
42
|
+
│ └── activations.py # tsallis_entmax (q-softmax / sparsemax family)
|
|
43
|
+
├── shared/
|
|
44
|
+
│ ├── __init__.py
|
|
45
|
+
│ ├── types.py # Array / Scalar type aliases
|
|
46
|
+
│ └── validation.py # q-range checks, broadcasting helpers
|
|
47
|
+
└── plots/
|
|
48
|
+
├── __init__.py
|
|
49
|
+
├── style.py # magma palette + rcParams + qcolors(n)
|
|
50
|
+
├── functions.py # plot q-log / q-exp families
|
|
51
|
+
└── distributions.py # plot q-Gaussian family
|
|
52
|
+
examples/ # runnable scripts (save figures with magma palette)
|
|
53
|
+
tests/ # pytest suite (limits, gradients, jit/vmap, shapes)
|
|
54
|
+
docs/ # theory note + API overview
|
|
55
|
+
pyproject.toml # uv-managed project (jax, matplotlib, pytest, ruff)
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
## 4. Mathematical scope (core)
|
|
59
|
+
|
|
60
|
+
| Primitive | Formula | `q → 1` limit |
|
|
61
|
+
|-----------|---------|---------------|
|
|
62
|
+
| `q_log(x, q)` | `(x**(1-q) - 1) / (1-q)` | `log(x)` |
|
|
63
|
+
| `q_exp(x, q)` | `[1 + (1-q)x]_+ ** (1/(1-q))` | `exp(x)` |
|
|
64
|
+
| `q_add(a, b, q)` | `a + b + (1-q) a b` | `a + b` |
|
|
65
|
+
| `tsallis_entropy(p, q)` | `(1 - Σ p**q) / (q-1)` | `-Σ p log p` |
|
|
66
|
+
| `tsallis_cross_entropy(p, y, q)` | `-Σ y · q_log(p, q)` | `-Σ y log p` |
|
|
67
|
+
| `tsallis_divergence(p, r, q)` | `(Σ p**q r**(1-q) - 1)/(q-1)` | `KL(p‖r)` |
|
|
68
|
+
| `q_gaussian_pdf(x, q, β)` | `√β/C_q · q_exp(-β x², q)` | `√(β/π)·exp(-βx²)` |
|
|
69
|
+
| `tsallis_entmax(z, q)` | argmax over simplex of `⟨p,z⟩ + Sᵀ_q(p)` | `softmax(z)` |
|
|
70
|
+
|
|
71
|
+
q-Gaussian sampling uses the **generalized Box–Muller** transform (valid for `q < 3`).
|
|
72
|
+
|
|
73
|
+
## 5. Enhancements over the bare skeleton
|
|
74
|
+
|
|
75
|
+
- Add `pyproject.toml` (uv) — the project currently has no build/dependency config.
|
|
76
|
+
- Add `shared/validation.py` for consistent `q` handling and clear errors.
|
|
77
|
+
- Centralize plot styling on **magma** (`plots/style.py`) instead of ad-hoc colors.
|
|
78
|
+
- Fill the six example stubs with runnable, figure-producing demos.
|
|
79
|
+
- Add a real pytest suite covering `q → 1` limits, gradient finiteness, `jit`/`vmap`.
|
|
80
|
+
- Add a concise theory note in `docs/`.
|
|
81
|
+
|
|
82
|
+
## 6. Checklist
|
|
83
|
+
|
|
84
|
+
### Project setup
|
|
85
|
+
- [x] `pyproject.toml` for uv (jax, matplotlib, pytest, ruff) + tool config
|
|
86
|
+
- [x] `README.md` with install + quickstart
|
|
87
|
+
- [x] `docs/theory.md` — concise Tsallis-statistics primer
|
|
88
|
+
|
|
89
|
+
### Core — math
|
|
90
|
+
- [x] `core/functions.py` — `q_log`, `q_exp`, `q_add`, `q_diff`, `q_prod`, `q_div`
|
|
91
|
+
- [x] `core/entropy.py` — `tsallis_entropy`, `tsallis_cross_entropy`, `tsallis_divergence`
|
|
92
|
+
- [x] `core/distributions.py` — q-Gaussian `pdf`, `logpdf`, `sample`, `normalization`
|
|
93
|
+
- [x] `core/activations.py` — `tsallis_entmax` (q-softmax / sparsemax family)
|
|
94
|
+
- [x] `core/__init__.py` — re-export core primitives
|
|
95
|
+
|
|
96
|
+
### Shared
|
|
97
|
+
- [x] `shared/types.py` — `Array`, `Scalar` aliases
|
|
98
|
+
- [x] `shared/validation.py` — `as_scalar_q`, range checks
|
|
99
|
+
- [x] `shared/__init__.py` — re-exports
|
|
100
|
+
|
|
101
|
+
### Plots (magma palette)
|
|
102
|
+
- [x] `plots/style.py` — `use_qjax_style()`, `qcolors(n)`, magma default
|
|
103
|
+
- [x] `plots/functions.py` — `plot_q_log`, `plot_q_exp`
|
|
104
|
+
- [x] `plots/distributions.py` — `plot_q_gaussian`
|
|
105
|
+
- [x] `plots/__init__.py` — re-exports
|
|
106
|
+
|
|
107
|
+
### Public API
|
|
108
|
+
- [x] `qjax/__init__.py` — curated exports + `__version__`
|
|
109
|
+
|
|
110
|
+
### Examples (runnable, save figures)
|
|
111
|
+
- [x] `examples/q_gaussian.py` — q-Gaussian family figure
|
|
112
|
+
- [x] `examples/learnable_q.py` — fit `q` by gradient descent to data
|
|
113
|
+
- [x] `examples/classification.py` — Tsallis cross-entropy classifier
|
|
114
|
+
- [x] `examples/optimization.py` — q-exp annealing / q-deformed objective
|
|
115
|
+
- [x] `examples/reinforcement_learning.py` — Tsallis-entmax policy on bandit
|
|
116
|
+
- [x] `examples/graph_learning.py` — Tsallis-entropy node attention on a graph
|
|
117
|
+
|
|
118
|
+
### Tests
|
|
119
|
+
- [x] `tests/test_functions.py` — limits, inverses, gradients, jit/vmap
|
|
120
|
+
- [x] `tests/test_entropy.py` — Shannon/KL limits, non-negativity
|
|
121
|
+
- [x] `tests/test_distributions.py` — normalization, sampling variance
|
|
122
|
+
- [x] `tests/test_activations.py` — simplex, softmax/sparsemax limits
|
|
123
|
+
|
|
124
|
+
### Verification
|
|
125
|
+
- [x] `uv run pytest` passes (69 tests)
|
|
126
|
+
- [x] `uv run ruff check` passes on `qjax`, `tests`, `examples`
|
|
127
|
+
- [x] Run every example end-to-end (6 figures saved to `examples/figures/`)
|
|
128
|
+
|
|
129
|
+
## 7. Notes from execution
|
|
130
|
+
|
|
131
|
+
- **Sampler correctness.** The first sampler used a generalized Box–Muller with an
|
|
132
|
+
incorrect scale (variance was ~20% off). Replaced with the exact Student-``t``
|
|
133
|
+
construction ``X = Z / sqrt(W/ν) / sqrt((3-q)β)``, which reproduces the family
|
|
134
|
+
variance ``1/((5-3q)β)`` exactly. Supported range is ``1 <= q < 3``.
|
|
135
|
+
- **Gradient bug in `normalization`.** The piecewise ``C_q`` closures captured the
|
|
136
|
+
raw ``q`` instead of their sanitized argument, so the unused branch evaluated
|
|
137
|
+
``log(1-q)`` with ``q > 1`` and poisoned the gradient via ``0 * NaN``. Fixed and
|
|
138
|
+
covered by `test_normalization_gradient_finite`. This is why `learnable_q` now
|
|
139
|
+
converges (recovers ``q ≈ 1.62`` from ``q_true = 1.6``).
|
|
140
|
+
- **`logpdf`.** Must use the ordinary ``log`` of the density, not ``q_log`` (which
|
|
141
|
+
would collapse to ``-βx²`` and drop the normalization curvature).
|
|
142
|
+
- All plots default to the **`magma`** colormap via `qjax.plots.use_qjax_style`.
|
|
143
|
+
|
|
144
|
+
## 8. Publication-grade revision
|
|
145
|
+
|
|
146
|
+
- **Vector PDF output.** `use_qjax_style` now configures a research-grade style
|
|
147
|
+
(serif body + Computer-Modern math via `mathtext`, embedded fonts with
|
|
148
|
+
`pdf.fonttype=42`, in-pointing major/minor ticks, top/right spines off, magma
|
|
149
|
+
color cycle). Added `qjax.plots.save_figure`, and every example saves a tight
|
|
150
|
+
vector **PDF** to `examples/figures/`.
|
|
151
|
+
- **Example rename.** `graph_learning.py` → `attention_mlp.py` (Tsallis-entmax
|
|
152
|
+
sparse attention).
|
|
153
|
+
- **Classification redesigned as a baseline study.** Now compares the **Shannon
|
|
154
|
+
baseline** (`q=1`, standard cross-entropy) against robust **Tsallis** losses
|
|
155
|
+
(`q=0.7`, `q=0.4`) across increasing label-noise levels, averaged over seeds on
|
|
156
|
+
a clean test set. Key correction: the *robust* Tsallis cross-entropy regime is
|
|
157
|
+
``q < 1`` (where ``-ln_q`` is bounded — the generalized cross-entropy), not
|
|
158
|
+
``q > 1``; all methods share softmax outputs and differ only in the loss `q`.
|
|
159
|
+
Result: clean-test accuracy degrades far more gracefully for `q<1`
|
|
160
|
+
(`q=0.4`: 0.84→0.68) than for Shannon (`q=1`: 0.82→0.59) as `η` rises to 0.5.
|
|
161
|
+
|
|
162
|
+
## 9. Documentation site
|
|
163
|
+
|
|
164
|
+
- **Sphinx + Furo.** `docs/` is a full MyST-Markdown documentation site:
|
|
165
|
+
`index`, `installation`, `quickstart`, `theory`, `examples`, and an autodoc
|
|
166
|
+
`api` reference (Napoleon Google-style docstrings, viewcode, intersphinx to
|
|
167
|
+
python/numpy/jax/matplotlib). Math via MyST `dollarmath` + MathJax; magma
|
|
168
|
+
accent colors matching the plot palette.
|
|
169
|
+
- Added a `docs` optional-dependency group (`sphinx`, `furo`, `myst-parser`) and
|
|
170
|
+
a `docs/Makefile`. Build: `uv run sphinx-build -b html docs docs/_build/html`
|
|
171
|
+
(clean under `-W`, warnings-as-errors).
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
# Read the Docs configuration file
|
|
2
|
+
# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
|
|
3
|
+
|
|
4
|
+
version: 2
|
|
5
|
+
|
|
6
|
+
build:
|
|
7
|
+
os: ubuntu-24.04
|
|
8
|
+
tools:
|
|
9
|
+
python: "3.13"
|
|
10
|
+
|
|
11
|
+
sphinx:
|
|
12
|
+
configuration: docs/conf.py
|
|
13
|
+
|
|
14
|
+
# python:
|
|
15
|
+
# install:
|
|
16
|
+
# - requirements: docs/requirements.txt
|
|
@@ -0,0 +1,131 @@
|
|
|
1
|
+
# Contributing to qjax
|
|
2
|
+
|
|
3
|
+
Thanks for your interest in `qjax`! It is a research library for **Tsallis
|
|
4
|
+
statistics in AI**, built on JAX. Contributions of all kinds are welcome: new
|
|
5
|
+
`q`-deformed primitives, examples, documentation, tests, and bug fixes.
|
|
6
|
+
|
|
7
|
+
This guide explains how to set up the project, the conventions we follow, and
|
|
8
|
+
how to get a change merged.
|
|
9
|
+
|
|
10
|
+
## Development setup
|
|
11
|
+
|
|
12
|
+
`qjax` is managed with [uv](https://docs.astral.sh/uv/) and targets
|
|
13
|
+
**Python 3.10+**.
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
git clone <repo-url> qjax
|
|
17
|
+
cd qjax
|
|
18
|
+
uv sync --extra dev # runtime + tests + linter
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
Useful commands (these mirror what CI runs):
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
uv run pytest # test suite
|
|
25
|
+
uv run pytest --cov=qjax --cov-report=term-missing # with coverage
|
|
26
|
+
uv run ruff check qjax tests examples # lint
|
|
27
|
+
uv run python examples/q_gaussian.py # run an example
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
To work on the documentation:
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
uv sync --extra docs
|
|
34
|
+
uv run sphinx-build -b html docs docs/_build/html -W # -W = warnings are errors
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Project layout
|
|
38
|
+
|
|
39
|
+
```
|
|
40
|
+
qjax/
|
|
41
|
+
├── core/ # the math: functions, entropy, distributions, activations
|
|
42
|
+
├── shared/ # type aliases and q validation helpers
|
|
43
|
+
└── plots/ # magma-themed, publication-grade plotting helpers
|
|
44
|
+
examples/ # runnable scripts that save figures to examples/figures/
|
|
45
|
+
tests/ # pytest suite
|
|
46
|
+
docs/ # Sphinx + Furo documentation (MyST Markdown)
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
`core/` is the single source of truth for the math; `plots/`, `examples/`, and
|
|
50
|
+
`tests/` only consume the public API.
|
|
51
|
+
|
|
52
|
+
## Design principles
|
|
53
|
+
|
|
54
|
+
New primitives should follow the conventions already in `core/`:
|
|
55
|
+
|
|
56
|
+
1. **Pure functions of `(x, q)`.** No hidden state. Every primitive must work
|
|
57
|
+
under `jax.jit`, `jax.grad`, and `jax.vmap`.
|
|
58
|
+
2. **Recover the classical limit.** As `q → 1`, the primitive must reduce to its
|
|
59
|
+
Boltzmann–Gibbs–Shannon counterpart (log, exp, Shannon entropy, KL, …).
|
|
60
|
+
3. **Finite gradients at `q = 1`.** Closed forms are typically `0/0` at `q = 1`.
|
|
61
|
+
Use the *double-`where`* trick: select the analytic limit with
|
|
62
|
+
`jnp.where(near_one(q), limit, deformed)`, and sanitize the unused branch's
|
|
63
|
+
inputs so it produces neither a `NaN` value nor a `NaN` gradient (a `0 * NaN`
|
|
64
|
+
from the unused branch will still poison `jax.grad`). See `qjax.core.functions`
|
|
65
|
+
and `qjax.core.distributions.normalization` for worked examples.
|
|
66
|
+
4. **Validate `q` consistently.** Use the helpers in `qjax.shared.validation`
|
|
67
|
+
(`as_scalar_q`, `near_one`, `Q_EPS`).
|
|
68
|
+
|
|
69
|
+
## Code style
|
|
70
|
+
|
|
71
|
+
- **Formatting/linting:** [ruff](https://github.com/astral-sh/ruff), line length
|
|
72
|
+
**100**. Run `uv run ruff check qjax tests examples` before committing; CI
|
|
73
|
+
enforces it.
|
|
74
|
+
- **Docstrings:** Google style (enforced by ruff's pydocstyle and rendered by
|
|
75
|
+
Sphinx/napoleon). Every public function documents its **formula**, its
|
|
76
|
+
**`q → 1` limit**, args (with shapes), and returns.
|
|
77
|
+
- **Type hints:** use the aliases in `qjax.shared.types` (`Array`, `Scalar`).
|
|
78
|
+
- Keep the public API curated in `qjax/__init__.py` and the relevant
|
|
79
|
+
`__init__.py` re-exports.
|
|
80
|
+
|
|
81
|
+
## Tests
|
|
82
|
+
|
|
83
|
+
Every new primitive needs tests in `tests/`. At minimum, cover:
|
|
84
|
+
|
|
85
|
+
- the **`q → 1` limit** (matches the classical function to tolerance);
|
|
86
|
+
- **gradient finiteness** with `jax.grad`, including at and near `q = 1`;
|
|
87
|
+
- **`jit` / `vmap`** compatibility and output shapes;
|
|
88
|
+
- any **mathematical identities** or domain edges (e.g. the Tsallis cut-off,
|
|
89
|
+
normalization, simplex constraints).
|
|
90
|
+
|
|
91
|
+
Tests run with float64 enabled (`jax.config.update("jax_enable_x64", True)`).
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
uv run pytest
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
## Plots and examples
|
|
98
|
+
|
|
99
|
+
- Use `qjax.plots.use_qjax_style()` and save with `qjax.plots.save_figure(...)`,
|
|
100
|
+
which writes vector **PDF** with the project's publication style.
|
|
101
|
+
- All figures use the **`magma`** colormap — get colors from
|
|
102
|
+
`qjax.plots.qcolors(n)` or the `CMAP` constant rather than hard-coding.
|
|
103
|
+
- Examples should be self-contained, reproducible (fixed PRNG seeds), and save
|
|
104
|
+
their output to `examples/figures/`.
|
|
105
|
+
|
|
106
|
+
## Submitting a change
|
|
107
|
+
|
|
108
|
+
1. Create a branch from `main`.
|
|
109
|
+
2. Make your change with tests and docstrings.
|
|
110
|
+
3. Ensure the full CI set passes locally:
|
|
111
|
+
```bash
|
|
112
|
+
uv run ruff check qjax tests examples
|
|
113
|
+
uv run pytest --cov=qjax --cov-report=term-missing
|
|
114
|
+
uv run sphinx-build -b html docs docs/_build/html -W
|
|
115
|
+
```
|
|
116
|
+
4. Open a pull request describing the change and the maths behind it. Link any
|
|
117
|
+
relevant references (the README and `docs/theory.md` cite the key papers).
|
|
118
|
+
|
|
119
|
+
CI (`.github/workflows/ci.yaml`) runs lint, the test matrix on Python
|
|
120
|
+
3.10–3.13, and the documentation build on every push and pull request.
|
|
121
|
+
|
|
122
|
+
## Reporting issues
|
|
123
|
+
|
|
124
|
+
When filing a bug, please include a **minimal reproducible example**, the
|
|
125
|
+
expected vs. actual behaviour, and your `jax` / `jaxlib` versions
|
|
126
|
+
(`uv run python -c "import jax; print(jax.__version__)"`).
|
|
127
|
+
|
|
128
|
+
## License
|
|
129
|
+
|
|
130
|
+
By contributing, you agree that your contributions are licensed under the
|
|
131
|
+
project's [MIT License](LICENSE).
|
qjax-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Kleyton da Costa
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
|
6
|
+
|
|
7
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
8
|
+
|
|
9
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
10
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
11
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
qjax-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: qjax
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Tsallis statistics for artificial intelligence, built on JAX.
|
|
5
|
+
Author-email: Kleyton Costa <kleyton.costa@holisticai.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Keywords: entropy,jax,machine-learning,statistics,tsallis
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Requires-Dist: jax>=0.4.30
|
|
11
|
+
Requires-Dist: matplotlib>=3.8
|
|
12
|
+
Provides-Extra: dev
|
|
13
|
+
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
|
|
14
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
15
|
+
Requires-Dist: ruff>=0.6; extra == 'dev'
|
|
16
|
+
Provides-Extra: docs
|
|
17
|
+
Requires-Dist: furo>=2024.1; extra == 'docs'
|
|
18
|
+
Requires-Dist: myst-parser>=2.0; extra == 'docs'
|
|
19
|
+
Requires-Dist: sphinx>=7.0; extra == 'docs'
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
<div align="center">
|
|
23
|
+
|
|
24
|
+
<img src="assets/logo_qjax.svg" alt="qjax logo" width="340"/>
|
|
25
|
+
|
|
26
|
+
# qjax
|
|
27
|
+
|
|
28
|
+
**Tsallis statistics for artificial intelligence, built on [JAX](https://github.com/jax-ml/jax).**
|
|
29
|
+
|
|
30
|
+
[](https://pypi.org/project/qjax/)
|
|
31
|
+
[](https://www.python.org/)
|
|
32
|
+
[](LICENSE)
|
|
33
|
+
[](https://github.com/jax-ml/jax)
|
|
34
|
+
[](https://github.com/astral-sh/ruff)
|
|
35
|
+
|
|
36
|
+
[Quickstart](#quickstart) • [Building blocks](#building-blocks) • [Example](#label-noise-robustness) • [Installation](#installation)
|
|
37
|
+
|
|
38
|
+
</div>
|
|
39
|
+
|
|
40
|
+
## What is qjax?
|
|
41
|
+
|
|
42
|
+
Tsallis (non-extensive) statistics generalizes Boltzmann–Gibbs–Shannon statistics through a single *entropic index* $q$. As $q \to 1$ every construction collapses back to its classical counterpart — Shannon entropy, the Gaussian, softmax, the Kullback–Leibler divergence — while $q \neq 1$ opens up heavy tails, sparse attention, and tunable exploration.
|
|
43
|
+
|
|
44
|
+
`qjax` exposes these $q$-deformed primitives as **pure, differentiable, `jit`/`vmap`-friendly** JAX functions. Because $q$ is just another argument, you can hold it fixed *or* **learn it end-to-end by gradient descent**.
|
|
45
|
+
|
|
46
|
+
Every primitive is a single closed form in the entropic index $q$, and each recovers its Boltzmann–Gibbs–Shannon counterpart in the $q \to 1$ limit:
|
|
47
|
+
|
|
48
|
+
| `qjax` | Definition | Limit $q \to 1$ |
|
|
49
|
+
| --- | --- | --- |
|
|
50
|
+
| `q_log` | $\ln_q x = \dfrac{x^{1-q} - 1}{1 - q}$ | $\ln x$ |
|
|
51
|
+
| `q_exp` | $\exp_q x = \big[1 + (1-q)\,x\big]_+^{\frac{1}{1-q}}$ | $e^{x}$ |
|
|
52
|
+
| `tsallis_entropy` | $S_q(p) = \dfrac{1 - \sum_i p_i^{\,q}}{q - 1}$ | $-\sum_i p_i \ln p_i$ |
|
|
53
|
+
| `tsallis_cross_entropy` | $H_q(y, p) = -\sum_i y_i \ln_q p_i$ | $-\sum_i y_i \ln p_i$ |
|
|
54
|
+
| `tsallis_divergence` | $D_q(p \,\Vert\, r) = \dfrac{\sum_i p_i^{\,q}\, r_i^{\,1-q} - 1}{q - 1}$ | $\mathrm{KL}(p \,\Vert\, r)$ |
|
|
55
|
+
| `q_gaussian_pdf` | $\mathcal{G}_q(x) = \dfrac{\sqrt{\beta}}{C_q}\,\exp_q(-\beta x^2)$ | $\sqrt{\tfrac{\beta}{\pi}}\,e^{-\beta x^2}$ |
|
|
56
|
+
| `tsallis_entmax` | $entmax_q(z) = \arg\max_{p \in \Delta}\,\langle p, z\rangle + S_q(p)$ | $softmax(z)$ |
|
|
57
|
+
|
|
58
|
+
where $[\,\cdot\,]_+ = \max(\cdot, 0)$ is the Tsallis cut-off, $C_q$ the $q$-Gaussian normalization, and $\Delta$ the probability simplex
|
|
59
|
+
(`tsallis_entmax` is exactly **sparsemax** at $q = 2$).
|
|
60
|
+
|
|
61
|
+
> `qjax` is a research library. The numerics are tested across the $q \to 1$
|
|
62
|
+
> limit, gradients, and `jit`/`vmap`, but the API may still evolve.
|
|
63
|
+
|
|
64
|
+
## Contents
|
|
65
|
+
|
|
66
|
+
- [Quickstart](#quickstart)
|
|
67
|
+
- [Building blocks](#building-blocks)
|
|
68
|
+
- [A learnable `q`](#a-learnable-q)
|
|
69
|
+
- [Label-noise robustness](#label-noise-robustness)
|
|
70
|
+
- [Installation](#installation)
|
|
71
|
+
- [Documentation](#documentation)
|
|
72
|
+
- [Development](#development)
|
|
73
|
+
- [Contributing](#contributing)
|
|
74
|
+
- [Citing](#citing)
|
|
75
|
+
- [License](#license)
|
|
76
|
+
|
|
77
|
+
## Quickstart
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
import jax, jax.numpy as jnp
|
|
81
|
+
import qjax
|
|
82
|
+
|
|
83
|
+
# q-deformed functions (recover log / exp as q -> 1)
|
|
84
|
+
qjax.q_log(2.0, q=1.5)
|
|
85
|
+
qjax.q_exp(1.0, q=1.5)
|
|
86
|
+
|
|
87
|
+
# Tsallis information measures
|
|
88
|
+
p = jnp.array([0.5, 0.3, 0.2])
|
|
89
|
+
qjax.tsallis_entropy(p, q=2.0) # -> Shannon entropy as q -> 1
|
|
90
|
+
qjax.tsallis_divergence(p, p, q=2.0) # -> KL divergence as q -> 1
|
|
91
|
+
|
|
92
|
+
# q-Gaussian distribution (heavy-tailed for 1 < q < 3)
|
|
93
|
+
x = jnp.linspace(-4, 4, 100)
|
|
94
|
+
qjax.q_gaussian_pdf(x, q=1.5, beta=1.0)
|
|
95
|
+
samples = qjax.sample(jax.random.PRNGKey(0), q=1.5, beta=1.0, shape=(1000,))
|
|
96
|
+
|
|
97
|
+
# Sparse softmax: q=1 -> softmax, q=2 -> sparsemax (exact zeros)
|
|
98
|
+
qjax.tsallis_entmax(jnp.array([2.0, 1.0, -1.0]), q=2.0)
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
## Building blocks
|
|
102
|
+
|
|
103
|
+
`qjax` is organized as a small set of composable, fully differentiable primitives. Each is a pure function of $(x, q)$.
|
|
104
|
+
|
|
105
|
+
### Deformed functions and $q$-algebra
|
|
106
|
+
|
|
107
|
+
`q_log` and `q_exp` are inverse deformations of `log`/`exp`; the accompanying $q$-algebra turns them into homomorphisms (`q_log(a·b) = q_add(q_log a, q_log b)`).
|
|
108
|
+
|
|
109
|
+
```python
|
|
110
|
+
qjax.q_log(x, q=1.5) # (x**(1-q) - 1) / (1-q)
|
|
111
|
+
qjax.q_add(qjax.q_log(2.0, 1.4), qjax.q_log(3.0, 1.4), 1.4) # == q_log(6.0, 1.4)
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
### Information measures
|
|
115
|
+
|
|
116
|
+
```python
|
|
117
|
+
p = jnp.array([0.5, 0.3, 0.2])
|
|
118
|
+
r = jnp.array([0.25, 0.25, 0.5])
|
|
119
|
+
|
|
120
|
+
qjax.tsallis_entropy(p, q=2.0) # -> Shannon entropy as q -> 1
|
|
121
|
+
qjax.tsallis_cross_entropy(p, r, q=2.0) # q-deformed cross-entropy loss
|
|
122
|
+
qjax.tsallis_divergence(p, r, q=2.0) # -> KL(p || r) as q -> 1
|
|
123
|
+
```
|
|
124
|
+
|
|
125
|
+
### The $q$-Gaussian
|
|
126
|
+
|
|
127
|
+
A maximum-Tsallis-entropy distribution: heavy-tailed (Student-$t$) for $1 < q < 3$, compactly supported for $q < 1$, Gaussian at $q = 1$.
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
x = jnp.linspace(-4, 4, 100)
|
|
131
|
+
qjax.q_gaussian_pdf(x, q=1.5, beta=1.0)
|
|
132
|
+
qjax.q_gaussian_logpdf(x, q=1.5, beta=1.0)
|
|
133
|
+
qjax.sample(jax.random.PRNGKey(0), q=1.5, beta=1.0, shape=(1000,))
|
|
134
|
+
```
|
|
135
|
+
|
|
136
|
+
### Sparse activations
|
|
137
|
+
|
|
138
|
+
`tsallis_entmax` interpolates between dense softmax ($q = 1$) and sparsemax
|
|
139
|
+
($q = 2$), producing exact zeros for $q > 1$ — a drop-in for sparse attention.
|
|
140
|
+
|
|
141
|
+
```python
|
|
142
|
+
z = jnp.array([2.0, 1.0, 0.1, -1.0])
|
|
143
|
+
qjax.tsallis_entmax(z, q=1.0) # softmax (dense)
|
|
144
|
+
qjax.tsallis_entmax(z, q=2.0) # sparsemax (exact zeros)
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
## A learnable $q$
|
|
148
|
+
|
|
149
|
+
Because $q$ is an ordinary differentiable argument, it is finite everywhere — including the $q = 1$ limit — so it can be optimized like any other parameter:
|
|
150
|
+
|
|
151
|
+
```python
|
|
152
|
+
import jax
|
|
153
|
+
|
|
154
|
+
x = jnp.linspace(-3, 3, 200)
|
|
155
|
+
nll = lambda q: -jnp.mean(qjax.q_gaussian_logpdf(x, q, 1.0))
|
|
156
|
+
grad_q = jax.grad(nll)(1.5) # well-defined gradient w.r.t. the entropic index
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
This is what makes $q$ more than a hyperparameter: the right amount of non-extensivity can be *discovered* from data.
|
|
160
|
+
|
|
161
|
+
## Label-noise robustness
|
|
162
|
+
|
|
163
|
+
When training labels are noisy, ordinary softmax **cross-entropy** is unbounded — a confidently mislabeled example incurs an arbitrarily large loss, so an over-parameterized network ends up *memorizing* the noise. Replacing the logarithm with the deformed $q$-logarithm gives the **Tsallis cross-entropy**, which is *bounded* for $q < 1$: its gradient saturates on unfittable points, so the model ignores label noise instead of fitting it.
|
|
164
|
+
|
|
165
|
+
For a one-hot target with true class $c$ and softmax probabilities $p$,
|
|
166
|
+
|
|
167
|
+
$$\mathcal{L}_q(p, c) = -\ln_q p_c = \frac{1 - p_c^{\,1-q}}{1 - q}, \qquad \ln_q x = \frac{x^{1-q} - 1}{1 - q}.$$
|
|
168
|
+
|
|
169
|
+
As $q \to 1$ this is exactly the standard cross-entropy $-\log p_c$; for $q < 1$ the per-example loss is bounded above by $1/(1-q)$, so mislabeled points cannot dominate the gradient.
|
|
170
|
+
|
|
171
|
+
The figure trains a small 3-class classifier on two shapes (blobs, spiral) from clean data up to 40% label noise, comparing the Boltzmann–Gibbs–Shannon baseline ($q = 1$) with Tsallis ($q = 0.3$). The comparison is fair — both share the same initialization, data, noisy labels and optimizer; only $q$ differs. Without noise the two match (≈98–99%); as noise grows the baseline carves spurious wrong-class islands while Tsallis keeps clean regions and higher accuracy.
|
|
172
|
+
|
|
173
|
+
<img src="docs/_static/examples/classification_boundaries.png" alt="Decision boundaries for blobs and spiral across noise levels: Tsallis vs the Boltzmann-Gibbs-Shannon baseline" width="960"/>
|
|
174
|
+
|
|
175
|
+
See the [classification example](docs/examples/classification.md) for the full setup.
|
|
176
|
+
|
|
177
|
+
## Installation
|
|
178
|
+
|
|
179
|
+
`qjax` requires Python 3.10+ and depends only on `jax` and `matplotlib`. It is managed with [uv](https://docs.astral.sh/uv/).
|
|
180
|
+
|
|
181
|
+
| Use case | Command |
|
|
182
|
+
| ------------------- | ---------------------------------------------------- |
|
|
183
|
+
| As a dependency | `uv add qjax` |
|
|
184
|
+
| Development | `uv sync --extra dev` (tests + linter) |
|
|
185
|
+
| Building the docs | `uv sync --extra docs` (Sphinx + Furo) |
|
|
186
|
+
|
|
187
|
+
For GPU/TPU acceleration, install the matching JAX build by following the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).
|
|
188
|
+
|
|
189
|
+
## Contributing
|
|
190
|
+
|
|
191
|
+
Contributions are welcome — new $q$-deformed primitives, examples, docs, and fixes. See [CONTRIBUTING.md](CONTRIBUTING.md) for the development setup, design principles (purity, the $q \to 1$ limit, finite gradients), and the checks CI runs.
|
|
192
|
+
|
|
193
|
+
## License
|
|
194
|
+
|
|
195
|
+
Released under the [MIT License](LICENSE).
|