API reference

The curated public API is re-exported at the top level (e.g. qjax.q_log); the canonical definitions live in the submodules documented below.

Core — deformed functions

q-deformed elementary functions and the Tsallis q-algebra.

This module implements the two foundational maps of non-extensive statistics, the q-logarithm and q-exponential, together with the q-deformed arithmetic they induce. Each function is a pure, differentiable JAX expression that recovers its Boltzmann–Gibbs counterpart as q -> 1.

Definitions

The q-logarithm and its inverse, the q-exponential, are

\[\ln_q(x) = \frac{x^{1-q} - 1}{1 - q}, \qquad \exp_q(x) = \big[1 + (1 - q)\,x\big]_+^{\frac{1}{1-q}},\]

where \([\cdot]_+ = \max(\cdot, 0)\). Both reduce to \(\ln\) and \(\exp\) as \(q \to 1\).

qjax.core.functions.q_log(x, q)[source]

q-logarithm \(\ln_q(x) = (x^{1-q} - 1)/(1-q)\).

Recovers the natural logarithm as q -> 1. The implementation uses the double-where trick so that the gradient is finite exactly at q = 1.

Parameters:
Returns:

Element-wise q-logarithm, same shape as x.

Return type:

Array

qjax.core.functions.q_exp(x, q)[source]

q-exponential \(\exp_q(x) = [1 + (1-q)x]_+^{1/(1-q)}\).

Inverse of q_log() and the limit of math.exp() as q -> 1. The base is clipped at zero (the Tsallis cut-off), so the result is 0 wherever 1 + (1-q)x < 0.

Parameters:
Returns:

Element-wise q-exponential, same shape as x.

Return type:

Array

qjax.core.functions.q_add(a, b, q)[source]

q-addition \(a \oplus_q b = a + b + (1-q)\,a\,b\).

The deformed sum satisfies q_log(x*y) = q_add(q_log(x), q_log(y)) and reduces to ordinary addition as q -> 1.

Parameters:
Returns:

Element-wise q-sum.

Return type:

Array

qjax.core.functions.q_diff(a, b, q)[source]

q-subtraction \(a \ominus_q b = (a - b)/(1 + (1-q)b)\).

Inverse of q_add() in its first argument: q_add(q_diff(a, b), b) == a.

Parameters:
Returns:

Element-wise q-difference.

Return type:

Array

qjax.core.functions.q_prod(a, b, q)[source]

q-product \(a \otimes_q b = [a^{1-q} + b^{1-q} - 1]_+^{1/(1-q)}\).

Dual to q_add(): it satisfies q_exp(x+y) = q_prod(q_exp(x), q_exp(y)) and reduces to ordinary multiplication as q -> 1. Defined for a, b >= 0.

Parameters:
Returns:

Element-wise q-product.

Return type:

Array

qjax.core.functions.q_div(a, b, q)[source]

q-division \(a \oslash_q b = [a^{1-q} - b^{1-q} + 1]_+^{1/(1-q)}\).

Inverse of q_prod() in its first argument and the limit of ordinary division as q -> 1. Defined for a, b >= 0.

Parameters:
Returns:

Element-wise q-quotient.

Return type:

Array

Core — entropy and divergences

Tsallis entropy, cross-entropy, and relative entropy (q-divergence).

These information measures generalize Shannon’s entropy and the Kullback–Leibler divergence through the entropic index q, recovering them in the limit q -> 1. They are the natural objective functions for non-extensive learning.

qjax.core.entropy.tsallis_entropy(p, q, axis=-1)[source]

Tsallis entropy \(S_q(p) = (1 - \sum_i p_i^q)/(q - 1)\).

Recovers the Shannon entropy \(-\sum_i p_i \log p_i\) as q -> 1. The entropy is concave in p and non-negative for probability vectors.

Parameters:

  • p (Array | float | int) – Probability mass values along axis (assumed non-negative and, ideally, normalized).

  • q (Array | float | int) – Entropic index (scalar).

  • axis (int) – Axis over which the distribution is defined.

Returns:

Tsallis entropy reduced over axis.

Return type:

Array

qjax.core.entropy.tsallis_cross_entropy(p, y, q, axis=-1)[source]

Tsallis cross-entropy \(H_q(y, p) = -\sum_i y_i \ln_q p_i\).

A drop-in q-deformed classification loss. With one-hot y it reduces to -q_log(p_correct, q), and to the standard cross-entropy as q -> 1.

Parameters:

  • p (Array | float | int) – Predicted probabilities along axis.

  • y (Array | float | int) – Target distribution along axis (e.g. one-hot labels).

  • q (Array | float | int) – Entropic index (scalar).

  • axis (int) – Axis over which the distributions are defined.

Returns:

Tsallis cross-entropy reduced over axis.

Return type:

Array

qjax.core.entropy.tsallis_divergence(p, r, q, axis=-1)[source]

Tsallis relative entropy \(D_q(p\,\|\,r)\).

Defined as \(D_q(p\|r) = \big(\sum_i p_i^q r_i^{1-q} - 1\big)/(q - 1)\), equivalently \(-\sum_i p_i \ln_q(r_i / p_i)\). Recovers the Kullback–Leibler divergence as q -> 1 and is non-negative.

Parameters:

  • p (Array | float | int) – First distribution along axis.

  • r (Array | float | int) – Second (reference) distribution along axis.

  • q (Array | float | int) – Entropic index (scalar).

  • axis (int) – Axis over which the distributions are defined.

Returns:

Tsallis divergence reduced over axis.

Return type:

Array

Core — the q-Gaussian distribution

The q-Gaussian distribution.

The q-Gaussian maximizes Tsallis entropy under a fixed second moment, just as the Gaussian maximizes Shannon entropy. Its density is

\[p(x) = \frac{\sqrt{\beta}}{C_q}\,\exp_q\!\big(-\beta x^2\big),\]

with normalization constant \(C_q\). It interpolates between heavy-tailed distributions (1 < q < 3; Student-t like) and compactly supported ones (q < 1), recovering the Gaussian as q -> 1.

qjax.core.distributions.normalization(q)[source]

Normalization constant \(C_q\) of the unit-\beta q-Gaussian.

Defined piecewise (see Tsallis, 2009) so that \(\int p(x)\,dx = 1\):

  • q < 1: \(C_q = \frac{2\sqrt{\pi}\,\Gamma(1/(1-q))} {(3-q)\sqrt{1-q}\,\Gamma(\tfrac{3-q}{2(1-q)})}\)

  • q = 1: \(C_q = \sqrt{\pi}\)

  • 1<q<3: \(C_q = \frac{\sqrt{\pi}\,\Gamma(\tfrac{3-q}{2(q-1)})} {\sqrt{q-1}\,\Gamma(1/(q-1))}\)

Parameters:

q (Array | float | int) – Entropic index (scalar), required to satisfy q < 3.

Returns:

The scalar normalization constant C_q.

Return type:

Array

qjax.core.distributions.q_gaussian_pdf(x, q=1.0, beta=1.0)[source]

Density of the q-Gaussian, \(\sqrt{\beta}/C_q\,\exp_q(-\beta x^2)\).

Parameters:
Returns:

Density values, same shape as x.

Return type:

Array

qjax.core.distributions.q_gaussian_logpdf(x, q=1.0, beta=1.0)[source]

Log-density of the q-Gaussian.

Computed as \(\tfrac12\log\beta - \log C_q + \log\!\exp_q(-\beta x^2)\). Returns -inf outside the (compact) support when q < 1, where the density vanishes.

Parameters:
Returns:

Log-density values, same shape as x.

Return type:

Array

qjax.core.distributions.sample(key, q=1.0, beta=1.0, shape=())[source]

Draw q-Gaussian samples for 1 <= q < 3.

For 1 < q < 3 the q-Gaussian is a rescaled Student-t: with \(\nu = (3-q)/(q-1)\) degrees of freedom,

\[X = \frac{T_\nu}{\sqrt{(3-q)\,\beta}}, \qquad T_\nu \sim \mathrm{t}(\nu),\]

where \(T_\nu = Z/\sqrt{W/\nu}\) with \(Z \sim \mathcal N(0,1)\) and \(W \sim \chi^2_\nu\). This yields exactly the family variance \(1/((5-3q)\beta)\) for q < 5/3. At q = 1 the Gaussian \(Z/\sqrt{2\beta}\) is returned.

Parameters:
Returns:

Samples of the requested shape.

Return type:

Array

Core — activations (entmax)

Tsallis entmax: the q-deformed softmax / sparsemax family.

entmax is the probability mapping obtained by regularizing the maximum-score problem with Tsallis entropy:

\[\mathrm{entmax}_q(z) = \arg\max_{p \in \Delta}\; \langle p, z \rangle + S_q^{T}(p),\]

with Tsallis entropy \(S_q^{T}(p) = \tfrac{1}{q(q-1)}(1 - \sum_i p_i^q)\). Its solution has the closed form (Peters, Niculae & Martins, 2019)

\[p_i = \big[(q - 1)\,z_i - \tau\big]_+^{\,1/(q-1)},\]

where the threshold \(\tau\) enforces \(\sum_i p_i = 1\). For q = 1 it is the ordinary softmax (dense); for q = 2 it is sparsemax (sparse). Intermediate q trade off smoothness against sparsity.

The threshold is found by bisection on a tight bracket, which is exact in the limit and fully compatible with jax.jit(), jax.grad(), and jax.vmap().

qjax.core.activations.tsallis_entmax(z, q=2.0, axis=-1, num_iters=50)[source]

Tsallis entmax over a simplex axis.

Solves for the threshold tau such that sum([(q-1)z - tau]_+^{1/(q-1)}) equals one, by bisection on the tight bracket [(q-1)·max(z) - 1, (q-1)·max(z)]. q = 1 short-circuits to a numerically stable softmax.

Parameters:

  • z (Array | float | int) – Scores / logits; the distribution is formed along axis.

  • q (Array | float | int) – Entropic index (scalar), q >= 1. q = 1 -> softmax, q = 2 -> sparsemax.

  • axis (int) – Axis over which to normalize.

  • num_iters (int) – Number of bisection steps for the threshold search.

Returns:

Probabilities with the same shape as z that sum to one along axis.

Return type:

Array

Shared — types and validation

Validation and broadcasting helpers for the entropic index q.

Tsallis primitives are parameterized by a single real number q (the entropic index). These helpers normalize q to a JAX scalar and provide a numerically robust mask for the q -> 1 limit, where most closed-form expressions become indeterminate (0 / 0).

qjax.shared.validation.Q_EPS: float = 1e-06

Absolute distance from q = 1 below which the Boltzmann–Gibbs limit is used.

qjax.shared.validation.as_scalar_q(q)[source]

Coerce an entropic index to a floating-point JAX scalar.

Parameters:

q (Array | float | int) – The entropic index, as a Python number or array-like.

Returns:

A 0-d jax.Array of floating dtype.

Return type:

Array

qjax.shared.validation.near_one(q, eps=1e-06)[source]

Boolean mask for indices that should use the q -> 1 (classical) limit.

Parameters:

  • q (Array | float | int) – The entropic index.

  • eps (float) – Half-width of the neighborhood around q = 1.

Returns:

A boolean array, broadcastable against q, that is True where |q - 1| < eps.

Return type:

Array

Plots

Publication-grade plotting style for qjax, themed on magma.

use_qjax_style() configures Matplotlib for research-grade, vector output: serif text with Computer-Modern math, embedded fonts, thin in-pointing ticks, and a magma color cycle. qcolors() samples a discrete sequence from magma so a family of curves indexed by q shares a coherent identity, and save_figure() writes a tight, font-embedded PDF.

qjax.plots.style.CMAP = 'magma'

The colormap used throughout qjax figures.

qjax.plots.style.qcolors(n, lo=0.1, hi=0.85)[source]

Sample n evenly spaced colors from the magma colormap.

The default [lo, hi] window avoids the near-black and near-white extremes so every curve stays legible on a white background.

Parameters:

  • n (int) – Number of colors to return.

  • lo (float) – Lower bound of the colormap window, in [0, 1].

  • hi (float) – Upper bound of the colormap window, in [0, 1].

Returns:

A list of n RGBA tuples.

Return type:

list

qjax.plots.style.use_qjax_style()[source]

Apply the qjax publication style (serif math, vector PDF, magma cycle).

Return type:

None

qjax.plots.style.save_figure(fig, path, transparent=False)[source]

Save fig as a tight, font-embedded vector PDF.

The extension is forced to .pdf and the parent directory is created if needed, so callers can pass a bare stem like figures/q_gaussian.

Parameters:

  • fig (Figure) – The figure to write.

  • path (str | Path) – Destination path; any extension is replaced with .pdf.

  • transparent (bool) – If True, write with a transparent background so the figure blends with whatever it is placed on.

Returns:

The resolved output path.

Return type:

Path

Plots of the q-deformed elementary functions across a range of q.

qjax.plots.functions.plot_q_log(q_values=(0.5, 0.8, 1.0, 1.5, 2.0), x_range=(0.05, 4.0), num=400, ax=None)[source]

Plot the q-logarithm for several entropic indices.

Parameters:

  • q_values (Sequence[float]) – Entropic indices to draw, one curve each.

  • x_range (tuple[float, float]) – (min, max) of the (positive) domain.

  • num (int) – Number of sample points.

  • ax (Axes | None) – Existing axis to draw on; a new one is created if None.

Returns:

The axis containing the plot.

Return type:

Axes

qjax.plots.functions.plot_q_exp(q_values=(0.5, 0.8, 1.0, 1.5, 2.0), x_range=(-3.0, 2.0), num=400, ax=None)[source]

Plot the q-exponential for several entropic indices.

Parameters:

  • q_values (Sequence[float]) – Entropic indices to draw, one curve each.

  • x_range (tuple[float, float]) – (min, max) of the domain.

  • num (int) – Number of sample points.

  • ax (Axes | None) – Existing axis to draw on; a new one is created if None.

Returns:

The axis containing the plot.

Return type:

Axes

Plots of the q-Gaussian density across a range of q.

qjax.plots.distributions.plot_q_gaussian(q_values=(0.5, 1.0, 1.5, 2.0, 2.5), beta=1.0, x_range=(-5.0, 5.0), num=500, ax=None)[source]

Plot the q-Gaussian density for several entropic indices.

Lower q gives compact support; q -> 1 is the Gaussian; higher q (up to 3) gives progressively heavier tails.

Parameters:

  • q_values (Sequence[float]) – Entropic indices to draw, one curve each (q < 3).

  • beta (float) – Shared inverse-width parameter.

  • x_range (tuple[float, float]) – (min, max) of the domain.

  • num (int) – Number of sample points.

  • ax (Axes | None) – Existing axis to draw on; a new one is created if None.

Returns:

The axis containing the plot.

Return type:

Axes