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.
Files changed (76) hide show
  1. qjax-0.1.0/.coverage +0 -0
  2. qjax-0.1.0/.github/workflows/ci.yaml +56 -0
  3. qjax-0.1.0/.github/workflows/release.yaml +27 -0
  4. qjax-0.1.0/.gitignore +15 -0
  5. qjax-0.1.0/.plan/PLAN.md +171 -0
  6. qjax-0.1.0/.readthedocs.yaml +16 -0
  7. qjax-0.1.0/CONTRIBUTING.md +131 -0
  8. qjax-0.1.0/LICENSE +11 -0
  9. qjax-0.1.0/PKG-INFO +195 -0
  10. qjax-0.1.0/README.md +174 -0
  11. qjax-0.1.0/assets/favicon.svg +32 -0
  12. qjax-0.1.0/assets/logo_qjax.svg +34 -0
  13. qjax-0.1.0/docs/Makefile +16 -0
  14. qjax-0.1.0/docs/_static/custom.css +11 -0
  15. qjax-0.1.0/docs/_static/examples/attention_mlp.png +0 -0
  16. qjax-0.1.0/docs/_static/examples/attention_q_learning.gif +0 -0
  17. qjax-0.1.0/docs/_static/examples/classification.png +0 -0
  18. qjax-0.1.0/docs/_static/examples/classification_boundaries.png +0 -0
  19. qjax-0.1.0/docs/_static/examples/learnable_q.png +0 -0
  20. qjax-0.1.0/docs/_static/examples/optimization.gif +0 -0
  21. qjax-0.1.0/docs/_static/examples/optimization.png +0 -0
  22. qjax-0.1.0/docs/_static/examples/q_gaussian.png +0 -0
  23. qjax-0.1.0/docs/_static/examples/reinforcement_learning.png +0 -0
  24. qjax-0.1.0/docs/_static/favicon.svg +32 -0
  25. qjax-0.1.0/docs/_static/logo.svg +34 -0
  26. qjax-0.1.0/docs/_templates/.gitkeep +0 -0
  27. qjax-0.1.0/docs/api.md +52 -0
  28. qjax-0.1.0/docs/conf.py +81 -0
  29. qjax-0.1.0/docs/examples/attention_mlp.md +64 -0
  30. qjax-0.1.0/docs/examples/attention_q_learning.md +52 -0
  31. qjax-0.1.0/docs/examples/classification.md +92 -0
  32. qjax-0.1.0/docs/examples/learnable_q.md +63 -0
  33. qjax-0.1.0/docs/examples/optimization.md +67 -0
  34. qjax-0.1.0/docs/examples/q_gaussian.md +68 -0
  35. qjax-0.1.0/docs/examples/reinforcement_learning.md +70 -0
  36. qjax-0.1.0/docs/examples.md +65 -0
  37. qjax-0.1.0/docs/index.md +144 -0
  38. qjax-0.1.0/docs/installation.md +64 -0
  39. qjax-0.1.0/docs/quickstart.md +67 -0
  40. qjax-0.1.0/docs/theory.md +104 -0
  41. qjax-0.1.0/examples/attention_mlp.py +281 -0
  42. qjax-0.1.0/examples/attention_q_learning.py +212 -0
  43. qjax-0.1.0/examples/classification.py +379 -0
  44. qjax-0.1.0/examples/figures/attention_mlp.pdf +0 -0
  45. qjax-0.1.0/examples/figures/attention_q_learning.gif +0 -0
  46. qjax-0.1.0/examples/figures/attention_q_learning.pdf +0 -0
  47. qjax-0.1.0/examples/figures/classification.pdf +0 -0
  48. qjax-0.1.0/examples/figures/classification_boundaries.pdf +0 -0
  49. qjax-0.1.0/examples/figures/learnable_q.pdf +0 -0
  50. qjax-0.1.0/examples/figures/optimization.gif +0 -0
  51. qjax-0.1.0/examples/figures/optimization.pdf +0 -0
  52. qjax-0.1.0/examples/figures/q_gaussian.pdf +0 -0
  53. qjax-0.1.0/examples/figures/reinforcement_learning.pdf +0 -0
  54. qjax-0.1.0/examples/learnable_q.py +69 -0
  55. qjax-0.1.0/examples/optimization.py +175 -0
  56. qjax-0.1.0/examples/q_gaussian.py +48 -0
  57. qjax-0.1.0/examples/reinforcement_learning.py +195 -0
  58. qjax-0.1.0/pyproject.toml +51 -0
  59. qjax-0.1.0/qjax/__init__.py +52 -0
  60. qjax-0.1.0/qjax/core/__init__.py +36 -0
  61. qjax-0.1.0/qjax/core/activations.py +111 -0
  62. qjax-0.1.0/qjax/core/distributions.py +155 -0
  63. qjax-0.1.0/qjax/core/entropy.py +108 -0
  64. qjax-0.1.0/qjax/core/functions.py +161 -0
  65. qjax-0.1.0/qjax/plots/__init__.py +15 -0
  66. qjax-0.1.0/qjax/plots/distributions.py +47 -0
  67. qjax-0.1.0/qjax/plots/functions.py +71 -0
  68. qjax-0.1.0/qjax/plots/style.py +116 -0
  69. qjax-0.1.0/qjax/shared/__init__.py +6 -0
  70. qjax-0.1.0/qjax/shared/types.py +20 -0
  71. qjax-0.1.0/qjax/shared/validation.py +45 -0
  72. qjax-0.1.0/tests/test_activations.py +69 -0
  73. qjax-0.1.0/tests/test_distributions.py +83 -0
  74. qjax-0.1.0/tests/test_entropy.py +111 -0
  75. qjax-0.1.0/tests/test_functions.py +71 -0
  76. 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
@@ -0,0 +1,15 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ .venv/
5
+ *.egg-info/
6
+
7
+ # Tooling caches
8
+ .pytest_cache/
9
+ .ruff_cache/
10
+
11
+ # Sphinx build output
12
+ docs/_build/
13
+
14
+ # Example outputs
15
+ examples/figures/*.png
@@ -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
+ [![PyPI](https://img.shields.io/pypi/v/qjax.svg)](https://pypi.org/project/qjax/)
31
+ [![Python](https://img.shields.io/badge/python-3.10%2B-3776ab.svg)](https://www.python.org/)
32
+ [![License: MIT](https://img.shields.io/badge/license-MIT-22863a.svg)](LICENSE)
33
+ [![Built on JAX](https://img.shields.io/badge/built%20on-JAX-b73779.svg)](https://github.com/jax-ml/jax)
34
+ [![Code style: Ruff](https://img.shields.io/badge/code%20style-ruff-261230.svg)](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).