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. Every entry on this page is rendered from the source docstrings by mkdocstrings, so it cannot disagree with the installed version.

How to read these pages

Signatures show the annotations from the source. Arguments, returns, and raises come from the Google-style docstring sections — the same text help() prints in a REPL. Click Source on any entry to see the implementation.

Core — deformed functions

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\).

Numerical form

Evaluating \((x^{1-q} - 1)/(1-q)\) directly loses catastrophically to cancellation as \(q \to 1\) — in float32 the relative error peaks near 4e-3 around q = 1.00001. Both functions are therefore written through jax.numpy.expm1 / jax.numpy.log1p, which are exact in that regime:

\[ \ln_q(x) = \log x \cdot \frac{e^{u} - 1}{u},\quad u = (1-q)\log x, \qquad \exp_q(x) = \exp\!\Big(x \cdot \frac{\log(1 + a)}{a}\Big),\quad a = (1-q)x. \]

The shared factor is the entire function \((e^t - 1)/t \to 1\) as \(t \to 0\), so no q = 1 special case is needed at all: the classical limit falls out of the same expression. Near \(t = 0\) a short Taylor series replaces the ratio, which keeps not only the value but also the derivative with respect to q correct — a hard where(q == 1, ...) branch would return a q-independent expression and hence a spurious zero q-gradient.

q_log

q_log(x: Array, q: Scalar) -> Array

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

Recovers the natural logarithm as q -> 1, continuously and with the correct derivative in q (see the module docstring).

At x = 0 the limit is -1/(1-q) for q < 1 and -inf for q >= 1. Negative x is outside the domain and yields NaN.

Parameters:

Name	Type	Description	Default
x	Array	Non-negative input, any shape.	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-logarithm, same shape as x.

Source code in qjax/core/functions.py

def q_log(x: Array, q: Scalar) -> jax.Array:
    r"""``q``-logarithm $\ln_q(x) = (x^{1-q} - 1)/(1-q)$.

    Recovers the natural logarithm as ``q -> 1``, continuously and with the
    correct derivative in ``q`` (see the module docstring).

    At ``x = 0`` the limit is ``-1/(1-q)`` for ``q < 1`` and ``-inf`` for
    ``q >= 1``. Negative ``x`` is outside the domain and yields ``NaN``.

    Args:
        x: Non-negative input, any shape.
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-logarithm, same shape as ``x``.
    """
    x = jnp.asarray(x, dtype=jnp.result_type(float))
    q = as_scalar_q(q)
    one_minus_q = 1.0 - q

    # Evaluate log on a sanitized argument so that x <= 0 contributes no NaN to
    # the gradient of the in-domain branch.
    positive = x > 0.0
    log_x = jnp.log(jnp.where(positive, x, 1.0))
    deformed = log_x * expm1_over_t(one_minus_q * log_x)

    safe_denom = jnp.where(one_minus_q == 0.0, 1.0, one_minus_q)
    at_zero = jnp.where(one_minus_q > 0.0, -1.0 / safe_denom, -jnp.inf)
    return jnp.where(positive, deformed, jnp.where(x == 0.0, at_zero, jnp.nan))

q_exp

q_exp(x: Array, q: Scalar) -> Array

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.

Past the Tsallis cut-off — that is, wherever 1 + (1-q)x <= 0 — the exponent 1/(1-q) decides the value: the result is 0 for q < 1 (positive exponent) and +inf for q > 1 (negative exponent), matching the genuine divergence of the q-exponential there.

Parameters:

Name	Type	Description	Default
x	Array	Input, any shape.	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-exponential, same shape as x.

Source code in qjax/core/functions.py

def q_exp(x: Array, q: Scalar) -> jax.Array:
    r"""``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``.

    Past the *Tsallis cut-off* — that is, wherever ``1 + (1-q)x <= 0`` — the
    exponent ``1/(1-q)`` decides the value: the result is ``0`` for ``q < 1``
    (positive exponent) and ``+inf`` for ``q > 1`` (negative exponent), matching
    the genuine divergence of the ``q``-exponential there.

    Args:
        x: Input, any shape.
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-exponential, same shape as ``x``.
    """
    x = jnp.asarray(x, dtype=jnp.result_type(float))
    q = as_scalar_q(q)
    one_minus_q = 1.0 - q
    a = one_minus_q * x

    # Sanitize before log1p so the clipped region contributes no NaN gradient.
    in_support = a > -1.0
    safe_a = jnp.where(in_support, a, 0.0)
    finite = jnp.exp(x * log1p_over_t(safe_a))

    cut_off = jnp.where(one_minus_q > 0.0, 0.0, jnp.inf)
    return jnp.where(in_support, finite, cut_off)

q_add

q_add(a: Array, b: Array, q: Scalar) -> Array

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:

Name	Type	Description	Default
a	Array	First operand.	required
b	Array	Second operand.	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-sum.

Source code in qjax/core/functions.py

def q_add(a: Array, b: Array, q: Scalar) -> jax.Array:
    r"""``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``.

    Args:
        a: First operand.
        b: Second operand.
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-sum.
    """
    a = jnp.asarray(a, dtype=jnp.result_type(float))
    b = jnp.asarray(b, dtype=jnp.result_type(float))
    one_minus_q = 1.0 - as_scalar_q(q)
    return a + b + one_minus_q * a * b

q_diff

q_diff(a: Array, b: Array, q: Scalar) -> Array

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:

Name	Type	Description	Default
a	Array	First operand.	required
b	Array	Second operand.	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-difference.

Source code in qjax/core/functions.py

def q_diff(a: Array, b: Array, q: Scalar) -> jax.Array:
    r"""``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``.

    Args:
        a: First operand.
        b: Second operand.
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-difference.
    """
    a = jnp.asarray(a, dtype=jnp.result_type(float))
    b = jnp.asarray(b, dtype=jnp.result_type(float))
    one_minus_q = 1.0 - as_scalar_q(q)
    return (a - b) / (1.0 + one_minus_q * b)

q_prod

q_prod(a: Array, b: Array, q: Scalar) -> Array

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:

Name	Type	Description	Default
a	Array	First operand (non-negative).	required
b	Array	Second operand (non-negative).	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-product.

Source code in qjax/core/functions.py

def q_prod(a: Array, b: Array, q: Scalar) -> jax.Array:
    r"""``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``.

    Args:
        a: First operand (non-negative).
        b: Second operand (non-negative).
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-product.
    """
    a = jnp.asarray(a, dtype=jnp.result_type(float))
    b = jnp.asarray(b, dtype=jnp.result_type(float))
    q = as_scalar_q(q)
    one_minus_q = 1.0 - q
    safe_exp = jnp.where(one_minus_q == 0.0, 1.0, one_minus_q)
    base = _safe_power(a, safe_exp) + _safe_power(b, safe_exp) - 1.0
    return jnp.where(near_one(q), a * b, _clipped_power(base, safe_exp))

q_div

q_div(a: Array, b: Array, q: Scalar) -> Array

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:

Name	Type	Description	Default
a	Array	Numerator (non-negative).	required
b	Array	Denominator (non-negative).	required
q	Scalar	Entropic index (scalar).	required

Returns:

Type	Description
Array	Element-wise q-quotient.

Source code in qjax/core/functions.py

def q_div(a: Array, b: Array, q: Scalar) -> jax.Array:
    r"""``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``.

    Args:
        a: Numerator (non-negative).
        b: Denominator (non-negative).
        q: Entropic index (scalar).

    Returns:
        Element-wise ``q``-quotient.
    """
    a = jnp.asarray(a, dtype=jnp.result_type(float))
    b = jnp.asarray(b, dtype=jnp.result_type(float))
    q = as_scalar_q(q)
    one_minus_q = 1.0 - q
    safe_exp = jnp.where(one_minus_q == 0.0, 1.0, one_minus_q)
    base = _safe_power(a, safe_exp) - _safe_power(b, safe_exp) + 1.0
    # Guard the classical quotient too: a bare ``a / b`` back-propagates NaN from
    # a zero denominator even when this branch is unselected (e.g. q = 2). The
    # divide-by-zero limit is spelled with constants rather than as ``a * inf``,
    # whose derivative is ``inf`` and would still poison the zero cotangent.
    nonzero_b = b != 0.0
    signed_inf = jnp.where(a > 0.0, jnp.inf, jnp.where(a < 0.0, -jnp.inf, jnp.nan))
    classical = jnp.where(nonzero_b, a / jnp.where(nonzero_b, b, 1.0), signed_inf)
    return jnp.where(near_one(q), classical, _clipped_power(base, safe_exp))

Core — entropy and divergences

entropy

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.

tsallis_entropy

tsallis_entropy(
    p: Array, q: Scalar, axis: int = -1
) -> Array

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.

Computed in the equivalent form \(\sum_i p_i \ln_q(1/p_i)\), which agrees with the definition above whenever p sums to one and, unlike it, stays finite and correctly differentiable in q at q = 1. The two forms differ on an unnormalized p, for which the definition above is genuinely singular at q = 1.

Parameters:

Name	Type	Description	Default
p	Array	Probability mass values along axis: non-negative and normalized.	required
q	Scalar	Entropic index (scalar).	required
axis	int	Axis over which the distribution is defined.	-1

Returns:

Type	Description
Array	Tsallis entropy reduced over axis.

Source code in qjax/core/entropy.py

def tsallis_entropy(p: Array, q: Scalar, axis: int = -1) -> jax.Array:
    r"""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.

    Computed in the equivalent form $\sum_i p_i \ln_q(1/p_i)$, which
    agrees with the definition above whenever ``p`` sums to one and, unlike it,
    stays finite and correctly differentiable in ``q`` at ``q = 1``. The two
    forms differ on an unnormalized ``p``, for which the definition above is
    genuinely singular at ``q = 1``.

    Args:
        p: Probability mass values along ``axis``: non-negative and normalized.
        q: Entropic index (scalar).
        axis: Axis over which the distribution is defined.

    Returns:
        Tsallis entropy reduced over ``axis``.
    """
    p = jnp.asarray(p, dtype=jnp.result_type(float))
    q = as_scalar_q(q)

    # Evaluated in the equivalent form S_q(p) = sum_i p_i ln_q(1/p_i), written
    # through the entire function (e^t - 1)/t. This is exact for a normalized p
    # and, unlike a hard `|q - 1| < eps` branch onto the Shannon expression,
    # keeps both the value and the derivative *with respect to q* correct
    # through q = 1 (the classical branch is q-independent, so its q-derivative
    # would be an identically zero cliff).
    #
    # The 0*log(0) = 0 convention is applied by masking on a sanitized p: a bare
    # p * log(p) is masked to 0 in value but still back-propagates
    # 0 * log(0) = 0 * -inf = NaN at a zero coordinate.
    support = p > 0.0
    safe_p = jnp.where(support, p, 1.0)
    log_p = jnp.log(safe_p)
    terms = jnp.where(support, -safe_p * log_p * expm1_over_t((q - 1.0) * log_p), 0.0)
    return jnp.sum(terms, axis=axis)

tsallis_cross_entropy

tsallis_cross_entropy(
    p: Array, y: Array, q: Scalar, axis: int = -1
) -> Array

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:

Name	Type	Description	Default
p	Array	Predicted probabilities along axis.	required
y	Array	Target distribution along axis (e.g. one-hot labels).	required
q	Scalar	Entropic index (scalar).	required
axis	int	Axis over which the distributions are defined.	-1

Returns:

Type	Description
Array	Tsallis cross-entropy reduced over axis.

Source code in qjax/core/entropy.py

def tsallis_cross_entropy(p: Array, y: Array, q: Scalar, axis: int = -1) -> jax.Array:
    r"""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``.

    Args:
        p: Predicted probabilities along ``axis``.
        y: Target distribution along ``axis`` (e.g. one-hot labels).
        q: Entropic index (scalar).
        axis: Axis over which the distributions are defined.

    Returns:
        Tsallis cross-entropy reduced over ``axis``.
    """
    p = jnp.asarray(p, dtype=jnp.result_type(float))
    y = jnp.asarray(y, dtype=jnp.result_type(float))
    # Apply the 0 * ln_q(0) = 0 convention. Where the target mass is zero the
    # term is dropped, and ``q_log`` is evaluated on a safe argument so a zero
    # prediction at an *unused* class (common for sparse ``tsallis_entmax``
    # outputs) yields neither a NaN value nor a NaN gradient. A zero prediction
    # *on* a positive-target class is a genuine +inf loss and is left intact.
    safe_p = jnp.where(y > 0.0, p, 1.0)
    contrib = jnp.where(y > 0.0, y * q_log(safe_p, q), 0.0)
    return -jnp.sum(contrib, axis=axis)

tsallis_divergence

tsallis_divergence(
    p: Array, r: Array, q: Scalar, axis: int = -1
) -> Array

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.

The second form is the one evaluated: it agrees with the first whenever p sums to one and, unlike it, stays correctly differentiable in q at q = 1.

Parameters:

Name	Type	Description	Default
p	Array	First distribution along axis, non-negative and normalized.	required
r	Array	Second (reference) distribution along axis.	required
q	Scalar	Entropic index (scalar).	required
axis	int	Axis over which the distributions are defined.	-1

Returns:

Type	Description
Array	Tsallis divergence reduced over axis.

Source code in qjax/core/entropy.py

def tsallis_divergence(p: Array, r: Array, q: Scalar, axis: int = -1) -> jax.Array:
    r"""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.

    The second form is the one evaluated: it agrees with the first whenever
    ``p`` sums to one and, unlike it, stays correctly differentiable in ``q`` at
    ``q = 1``.

    Args:
        p: First distribution along ``axis``, non-negative and normalized.
        r: Second (reference) distribution along ``axis``.
        q: Entropic index (scalar).
        axis: Axis over which the distributions are defined.

    Returns:
        Tsallis divergence reduced over ``axis``.
    """
    p = jnp.asarray(p, dtype=jnp.result_type(float))
    r = jnp.asarray(r, dtype=jnp.result_type(float))
    q = as_scalar_q(q)
    q_minus_one = q - 1.0

    # Evaluated as D_q(p||r) = -sum_i p_i ln_q(r_i / p_i) — the equivalent form
    # already named in the docstring — written through the entire function
    # (e^t - 1)/t so that the Kullback-Leibler limit, and its q-derivative, come
    # out of the same expression instead of a separate q-independent branch.
    #
    # One mask serves both regimes. An earlier version guarded only the KL
    # branch, which still leaked NaN: the deformed branch evaluates
    # r ** (1-q) = 0 ** negative = inf at a zero reference and back-propagates
    # NaN even when unselected.
    mask = (p > 0.0) & (r > 0.0)
    safe_p = jnp.where(mask, p, 1.0)
    safe_r = jnp.where(mask, r, 1.0)
    log_ratio = jnp.log(safe_p / safe_r)
    supported = safe_p * log_ratio * expm1_over_t(q_minus_one * log_ratio)

    # Where p > 0 but r == 0 the reference assigns no mass to an outcome p deems
    # possible: ln_q(0) is -inf for q >= 1, giving a genuinely divergent term,
    # and -1/(1-q) for q < 1, giving the finite p/(1-q). Where p == 0 the term
    # vanishes under the 0 * ln_q(0) = 0 convention.
    safe_gap = jnp.where(q_minus_one < 0.0, -q_minus_one, 1.0)
    unsupported = jnp.where(q_minus_one < 0.0, p / safe_gap, jnp.inf)
    terms = jnp.where(mask, supported, jnp.where(p > 0.0, unsupported, 0.0))
    return jnp.sum(terms, axis=axis)

Core — the q-Gaussian distribution

distributions

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.

normalization

normalization(q: Scalar) -> Array

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:

Name	Type	Description	Default
q	Scalar	Entropic index (scalar), required to satisfy q < 3.	required

Returns:

Type	Description
Array	The scalar normalization constant C_q.

Source code in qjax/core/distributions.py

def normalization(q: Scalar) -> jax.Array:
    r"""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))}$

    Args:
        q: Entropic index (scalar), required to satisfy ``q < 3``.

    Returns:
        The scalar normalization constant ``C_q``.
    """
    q = as_scalar_q(q)
    sqrt_pi = jnp.sqrt(jnp.pi)

    def below(qb):  # q < 1
        a = 1.0 - qb
        log_c = (
            jnp.log(2.0)
            + 0.5 * jnp.log(jnp.pi)
            + gammaln(1.0 / a)
            - jnp.log(3.0 - qb)
            - 0.5 * jnp.log(a)
            - gammaln((3.0 - qb) / (2.0 * a))
        )
        return jnp.exp(log_c)

    def above(qa):  # 1 < q < 3
        b = qa - 1.0
        log_c = (
            0.5 * jnp.log(jnp.pi)
            + gammaln((3.0 - qa) / (2.0 * b))
            - 0.5 * jnp.log(b)
            - gammaln(1.0 / b)
        )
        return jnp.exp(log_c)

    # Each branch is evaluated with a sanitized argument that stays strictly
    # inside its own domain, so the unused branch produces neither a NaN value
    # nor a NaN gradient before the final selection by the sign of (q - 1).
    q_safe_below = jnp.where(q < 1.0, q, 0.0)
    q_safe_above = jnp.where(q > 1.0, q, 2.0)
    c_below = below(q_safe_below)
    c_above = above(q_safe_above)
    return jnp.where(q < 1.0, c_below, jnp.where(q > 1.0, c_above, sqrt_pi))

q_gaussian_pdf

q_gaussian_pdf(
    x: Array, q: Scalar = 1.0, beta: Scalar = 1.0
) -> Array

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

Parameters:

Name	Type	Description	Default
x	Array	Evaluation points, any shape.	required
q	Scalar	Entropic index (scalar), q < 3.	1.0
beta	Scalar	Inverse-width parameter beta > 0.	1.0

Returns:

Type	Description
Array	Density values, same shape as x.

Source code in qjax/core/distributions.py

def q_gaussian_pdf(x: Array, q: Scalar = 1.0, beta: Scalar = 1.0) -> jax.Array:
    r"""Density of the ``q``-Gaussian, $\sqrt{\beta}/C_q\,\exp_q(-\beta x^2)$.

    Args:
        x: Evaluation points, any shape.
        q: Entropic index (scalar), ``q < 3``.
        beta: Inverse-width parameter ``beta > 0``.

    Returns:
        Density values, same shape as ``x``.
    """
    x = jnp.asarray(x, dtype=jnp.result_type(float))
    beta = jnp.asarray(beta, dtype=jnp.result_type(float))
    return jnp.sqrt(beta) / normalization(q) * q_exp(-beta * x**2, q)

q_gaussian_logpdf

q_gaussian_logpdf(
    x: Array, q: Scalar = 1.0, beta: Scalar = 1.0
) -> Array

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:

Name	Type	Description	Default
x	Array	Evaluation points, any shape.	required
q	Scalar	Entropic index (scalar), q < 3.	1.0
beta	Scalar	Inverse-width parameter beta > 0.	1.0

Returns:

Type	Description
Array	Log-density values, same shape as x.

Source code in qjax/core/distributions.py

def q_gaussian_logpdf(x: Array, q: Scalar = 1.0, beta: Scalar = 1.0) -> jax.Array:
    r"""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.

    Args:
        x: Evaluation points, any shape.
        q: Entropic index (scalar), ``q < 3``.
        beta: Inverse-width parameter ``beta > 0``.

    Returns:
        Log-density values, same shape as ``x``.
    """
    x = jnp.asarray(x, dtype=jnp.result_type(float))
    beta = jnp.asarray(beta, dtype=jnp.result_type(float))
    log_prefactor = 0.5 * jnp.log(beta) - jnp.log(normalization(q))
    return log_prefactor + jnp.log(q_exp(-beta * x**2, q))

sample

sample(
    key: Array,
    q: Scalar = 1.0,
    beta: Scalar = 1.0,
    shape: tuple[int, ...] = (),
) -> Array

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:

Name	Type	Description	Default
key	Array	A jax.random.PRNGKey.	required
q	Scalar	Entropic index (scalar), 1 <= q < 3.	1.0
beta	Scalar	Inverse-width parameter beta > 0.	1.0
shape	tuple[int, ...]	Output shape.	()

Returns:

Type	Description
Array	Samples of the requested shape.

Source code in qjax/core/distributions.py

def sample(
    key: jax.Array,
    q: Scalar = 1.0,
    beta: Scalar = 1.0,
    shape: tuple[int, ...] = (),
) -> jax.Array:
    r"""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.

    Args:
        key: A `jax.random.PRNGKey`.
        q: Entropic index (scalar), ``1 <= q < 3``.
        beta: Inverse-width parameter ``beta > 0``.
        shape: Output shape.

    Returns:
        Samples of the requested ``shape``.
    """
    q = as_scalar_q(q)
    beta = jnp.asarray(beta, dtype=jnp.result_type(float))
    k_z, k_w = jax.random.split(key)
    z = jax.random.normal(k_z, shape)

    def gaussian(_):
        return z / jnp.sqrt(2.0 * beta)

    def student_t(_):
        nu = (3.0 - q) / (q - 1.0)
        w = 2.0 * jax.random.gamma(k_w, nu / 2.0, shape)  # chi-square with nu dof
        t = z / jnp.sqrt(w / nu)
        return t / jnp.sqrt((3.0 - q) * beta)

    return jax.lax.cond(near_one(q), gaussian, student_t, operand=None)

Core — activations (entmax)

activations

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; q = 2 is sparsemax. Larger q gives sparser distributions, q < 1 gives distributions denser (higher entropy) than softmax, and q -> 0^+ approaches the uniform distribution.

Differentiation

The threshold is located by bisection, but the mapping is not differentiated through that loop — doing so yields a wrong (and non-symmetric) Jacobian. The solve is wrapped in jax.lax.stop_gradient and the derivative is supplied by a jax.custom_jvp rule built from the implicit function theorem.

Writing \(s_i = p_i^{2-q}\) and \(h_i = -p_i \log p_i\) (both zero off the support), and letting \(T(v) = v - s\,\sum_j v_j / \sum_j s_j\), both tangents share a single form:

\[ \dot p = T\!\Big(s \odot \dot z + \big(h + s \odot z\big)\tfrac{\dot q}{q - 1}\Big), \]

so the Jacobian with respect to z is \(J = \mathrm{diag}(s) - s s^\top / \sum_i s_i\) — symmetric, positive semi-definite, and annihilating the all-ones vector (entmax is invariant to a constant shift of z). Because custom_jvp is used rather than custom_vjp, forward mode, reverse mode, and higher-order derivatives are all exact.

tsallis_entmax

tsallis_entmax(
    z: Array,
    q: Scalar = 2.0,
    axis: int = -1,
    num_iters: int = 50,
) -> Array

Tsallis entmax over a simplex axis.

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

Gradients are exact: the threshold search is not differentiated through, and a jax.custom_jvp rule supplies the implicit-function derivative with respect to both z and q. The Jacobian in z is diag(s) - s s^T / sum(s) with s = p ** (2 - q).

Parameters:

Name	Type	Description	Default
z	Array	Scores / logits; the distribution is formed along axis.	required
q	Scalar	Entropic index (scalar), q > 0. q = 1 -> softmax, q = 2 -> sparsemax. Values above 1 give sparse outputs; values below 1 give outputs denser than softmax.	2.0
axis	int	Axis over which to normalize.	-1
num_iters	int	Number of bisection steps for the threshold search. A Newton polish step follows, so the result is accurate even for small values.	50

Returns:

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

Raises:

Type	Description
ValueError	If q is a concrete value and q <= 0. A traced, non-positive q yields NaN instead.

Source code in qjax/core/activations.py

def tsallis_entmax(
    z: Array,
    q: Scalar = 2.0,
    axis: int = -1,
    num_iters: int = 50,
) -> jax.Array:
    r"""Tsallis ``entmax`` over a simplex axis.

    Solves for the threshold ``tau`` such that ``sum([(q-1)z - tau]_+^{1/(q-1)})``
    equals one. ``q = 1`` short-circuits to a numerically stable softmax.

    Gradients are exact: the threshold search is not differentiated through, and
    a `jax.custom_jvp` rule supplies the implicit-function derivative with
    respect to both ``z`` and ``q``. The Jacobian in ``z`` is
    ``diag(s) - s s^T / sum(s)`` with ``s = p ** (2 - q)``.

    Args:
        z: Scores / logits; the distribution is formed along ``axis``.
        q: Entropic index (scalar), ``q > 0``. ``q = 1`` -> softmax,
            ``q = 2`` -> sparsemax. Values above ``1`` give sparse outputs;
            values below ``1`` give outputs denser than softmax.
        axis: Axis over which to normalize.
        num_iters: Number of bisection steps for the threshold search. A Newton
            polish step follows, so the result is accurate even for small values.

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

    Raises:
        ValueError: If ``q`` is a concrete value and ``q <= 0``. A traced,
            non-positive ``q`` yields ``NaN`` instead.
    """
    z = jnp.asarray(z, dtype=jnp.result_type(float))
    q = positive_q_or_nan(as_scalar_q(q))

    # Move the working axis to the end for a uniform reduction layout.
    z = jnp.moveaxis(z, axis, -1)

    near = near_one(q)
    # Under vmap over a batched q, lax.cond lowers to a select and *both*
    # branches execute, so keep the entmax branch clear of the 1/(q-1) pole.
    # The clamped region is exactly the region the select discards.
    q_eff = jnp.where(near, jnp.where(q < 1.0, 1.0 - Q_EPS, 1.0 + Q_EPS), q)

    p = jax.lax.cond(
        near,
        lambda _: _softmax_core(z, q),
        lambda _: _entmax_core(z, q_eff, num_iters),
        operand=None,
    )
    return jnp.moveaxis(p, -1, axis)

Neural-network building blocks

Framework-agnostic pieces for Tsallis models: everything here operates on plain arrays and pytrees, so it composes with Flax, Equinox, Haiku, or hand-rolled JAX without adding a dependency on any of them.

reparam

Keeping a learnable entropic index inside its valid range.

Treating q as a free parameter and optimizing it directly does not work: the Tsallis primitives are undefined at q <= 0, the q-Gaussian requires q < 3, and gradient descent will happily step outside either bound. The standard remedy — used verbatim in six of this repository's examples before it lived here — is to optimize an unconstrained real q_raw and squash it:

\[ q = q_{\min} + (q_{\max} - q_{\min})\,\sigma(q_{\mathrm{raw}}). \]

bounded_q

bounded_q(
    q_raw: Scalar, lo: Scalar = 1.0, hi: Scalar = 3.0
) -> Array

Map an unconstrained parameter to an entropic index in (lo, hi).

The map is smooth and strictly monotone, so gradients flow to q_raw everywhere and the optimizer can never leave the interval. q_raw = 0 corresponds to the midpoint.

The interval is open in exact arithmetic but closed in floating point: the sigmoid saturates to exactly 0 or 1 once |q_raw| exceeds roughly 37 (float64) or 17 (float32), so q can reach lo or hi exactly. Choose lo strictly inside the valid domain — q > 0 for the Tsallis primitives, q < 3 for the q-Gaussian — rather than relying on strict inequality here.

Parameters:

Name	Type	Description	Default
q_raw	Scalar	Unconstrained real parameter, any shape.	required
lo	Scalar	Lower bound of the open interval (exclusive).	1.0
hi	Scalar	Upper bound of the open interval (exclusive).	3.0

Returns:

Type	Description
Array	The entropic index, same shape as q_raw.

Example

import jax.numpy as jnp from qjax.nn import bounded_q float(bounded_q(jnp.asarray(0.0), 1.0, 3.0)) 2.0

Source code in qjax/nn/reparam.py

def bounded_q(q_raw: Scalar, lo: Scalar = 1.0, hi: Scalar = 3.0) -> jax.Array:
    r"""Map an unconstrained parameter to an entropic index in ``(lo, hi)``.

    The map is smooth and strictly monotone, so gradients flow to ``q_raw``
    everywhere and the optimizer can never leave the interval. ``q_raw = 0``
    corresponds to the midpoint.

    The interval is open in exact arithmetic but closed in floating point: the
    sigmoid saturates to exactly ``0`` or ``1`` once ``|q_raw|`` exceeds roughly
    ``37`` (float64) or ``17`` (float32), so ``q`` can reach ``lo`` or ``hi``
    exactly. Choose ``lo`` strictly inside the valid domain — ``q > 0`` for the
    Tsallis primitives, ``q < 3`` for the ``q``-Gaussian — rather than relying
    on strict inequality here.

    Args:
        q_raw: Unconstrained real parameter, any shape.
        lo: Lower bound of the open interval (exclusive).
        hi: Upper bound of the open interval (exclusive).

    Returns:
        The entropic index, same shape as ``q_raw``.

    Example:
        >>> import jax.numpy as jnp
        >>> from qjax.nn import bounded_q
        >>> float(bounded_q(jnp.asarray(0.0), 1.0, 3.0))
        2.0
    """
    lo = jnp.asarray(lo, dtype=jnp.result_type(float))
    hi = jnp.asarray(hi, dtype=jnp.result_type(float))
    return lo + (hi - lo) * jax.nn.sigmoid(jnp.asarray(q_raw, dtype=jnp.result_type(float)))

inverse_bounded_q

inverse_bounded_q(
    q: Scalar, lo: Scalar = 1.0, hi: Scalar = 3.0
) -> Array

Invert bounded_q to initialize q_raw at a chosen q.

Useful for starting training from a meaningful index — q = 1 for a softmax-like attention map, say — rather than from the interval midpoint.

Parameters:

Name	Type	Description	Default
q	Scalar	Target entropic index, strictly inside (lo, hi).	required
lo	Scalar	Lower bound used by bounded_q.	1.0
hi	Scalar	Upper bound used by bounded_q.	3.0

Returns:

Type	Description
Array	The q_raw for which bounded_q(q_raw, lo, hi) == q.

Source code in qjax/nn/reparam.py

def inverse_bounded_q(q: Scalar, lo: Scalar = 1.0, hi: Scalar = 3.0) -> jax.Array:
    """Invert `bounded_q` to initialize ``q_raw`` at a chosen ``q``.

    Useful for starting training from a meaningful index — ``q = 1`` for a
    softmax-like attention map, say — rather than from the interval midpoint.

    Args:
        q: Target entropic index, strictly inside ``(lo, hi)``.
        lo: Lower bound used by `bounded_q`.
        hi: Upper bound used by `bounded_q`.

    Returns:
        The ``q_raw`` for which ``bounded_q(q_raw, lo, hi) == q``.
    """
    q = jnp.asarray(q, dtype=jnp.result_type(float))
    lo = jnp.asarray(lo, dtype=jnp.result_type(float))
    hi = jnp.asarray(hi, dtype=jnp.result_type(float))
    unit = (q - lo) / (hi - lo)
    return jnp.log(unit) - jnp.log1p(-unit)

attention

Attention whose normalizer is tsallis_entmax.

Replacing the softmax in an attention block with entmax makes the sparsity of the attention map a property of the entropic index rather than a fixed choice: q = 1 reproduces ordinary softmax attention, q = 2 gives sparsemax (most positions receive exactly zero weight), and q < 1 spreads the weight more evenly than softmax. Because q is differentiable, it can be learned jointly with the rest of the network — see bounded_q.

entmax_attention

entmax_attention(
    queries: Array,
    keys: Array,
    values: Array,
    q: Scalar = 2.0,
    mask: Array | None = None,
    scale: Scalar | None = None,
    num_iters: int = 50,
) -> tuple[Array, Array]

Scaled dot-product attention normalized by tsallis_entmax.

Computes entmax_q(Q K^T / sqrt(d)) V over the last (key) axis. Leading axes broadcast, so the same call serves single-head, multi-head, and batched inputs.

Parameters:

Name	Type	Description	Default
queries	Array	Query vectors, shape (..., n_queries, d_key). A single query per batch element may be passed as (..., d_key).	required
keys	Array	Key vectors, shape (..., n_keys, d_key).	required
values	Array	Value vectors, shape (..., n_keys, d_value).	required
q	Scalar	Entropic index (scalar), q > 0. 1 is softmax attention, 2 is sparsemax attention.	2.0
mask	Array | None	Optional boolean array broadcastable to the score shape (..., n_queries, n_keys). False positions are excluded.	None
scale	Scalar | None	Divisor applied to the scores. Defaults to sqrt(d_key).	None
num_iters	int	Bisection steps for the entmax threshold search.	50

Returns:

Type	Description
Array	A (context, attention) pair. context has shape
Array	(..., n_queries, d_value) and attention has shape
tuple[Array, Array]	(..., n_queries, n_keys) and sums to one over the last axis.

Example

import jax.numpy as jnp from qjax.nn import entmax_attention q_vec = jnp.ones((2, 4)) k = jnp.ones((2, 5, 4)) v = jnp.ones((2, 5, 3)) context, attn = entmax_attention(q_vec, k, v, q=2.0) context.shape, attn.shape ((2, 3), (2, 5))

Source code in qjax/nn/attention.py

def entmax_attention(
    queries: Array,
    keys: Array,
    values: Array,
    q: Scalar = 2.0,
    mask: Array | None = None,
    scale: Scalar | None = None,
    num_iters: int = 50,
) -> tuple[jax.Array, jax.Array]:
    r"""Scaled dot-product attention normalized by `tsallis_entmax`.

    Computes ``entmax_q(Q K^T / sqrt(d)) V`` over the last (key) axis. Leading
    axes broadcast, so the same call serves single-head, multi-head, and batched
    inputs.

    Args:
        queries: Query vectors, shape ``(..., n_queries, d_key)``. A single
            query per batch element may be passed as ``(..., d_key)``.
        keys: Key vectors, shape ``(..., n_keys, d_key)``.
        values: Value vectors, shape ``(..., n_keys, d_value)``.
        q: Entropic index (scalar), ``q > 0``. ``1`` is softmax attention,
            ``2`` is sparsemax attention.
        mask: Optional boolean array broadcastable to the score shape
            ``(..., n_queries, n_keys)``. ``False`` positions are excluded.
        scale: Divisor applied to the scores. Defaults to ``sqrt(d_key)``.
        num_iters: Bisection steps for the ``entmax`` threshold search.

    Returns:
        A ``(context, attention)`` pair. ``context`` has shape
        ``(..., n_queries, d_value)`` and ``attention`` has shape
        ``(..., n_queries, n_keys)`` and sums to one over the last axis.

    Example:
        >>> import jax.numpy as jnp
        >>> from qjax.nn import entmax_attention
        >>> q_vec = jnp.ones((2, 4))
        >>> k = jnp.ones((2, 5, 4))
        >>> v = jnp.ones((2, 5, 3))
        >>> context, attn = entmax_attention(q_vec, k, v, q=2.0)
        >>> context.shape, attn.shape
        ((2, 3), (2, 5))
    """
    queries = jnp.asarray(queries, dtype=jnp.result_type(float))
    keys = jnp.asarray(keys, dtype=jnp.result_type(float))
    values = jnp.asarray(values, dtype=jnp.result_type(float))

    # A single query per batch element is the common case in attention pooling;
    # add the query axis, then drop it again on the way out.
    squeeze_query = queries.ndim == keys.ndim - 1
    if squeeze_query:
        queries = queries[..., None, :]

    if scale is None:
        scale = jnp.sqrt(jnp.asarray(keys.shape[-1], dtype=queries.dtype))
    scores = jnp.einsum("...qd,...kd->...qk", queries, keys) / scale

    if mask is not None:
        # -inf scores are driven to exactly zero weight by entmax for every q,
        # and keep the masked positions out of the threshold search.
        scores = jnp.where(jnp.asarray(mask, dtype=bool), scores, -jnp.inf)

    attention = tsallis_entmax(scores, q=q, axis=-1, num_iters=num_iters)
    context = jnp.einsum("...qk,...kd->...qd", attention, values)

    if squeeze_query:
        context = context[..., 0, :]
        attention = attention[..., 0, :]
    return context, attention

losses

Loss functions built on the Tsallis information measures.

The Tsallis cross-entropy \(H_q(y, p) = -\sum_i y_i \ln_q p_i\) is a drop-in replacement for the usual cross-entropy that recovers it at q = 1.

Its practical appeal is robustness to label noise, and that lives at q < 1. There the q-logarithm is bounded below,

\[ \ln_q(0) = \frac{-1}{1 - q} \quad (q < 1), \]

so a confidently wrong prediction — the signature of a mislabelled example -- contributes at most \(1/(1-q)\) instead of diverging, and its gradient is capped with it. At q = 1 the penalty is unbounded, and for q > 1 it grows faster than the logarithm (like \(p^{1-q}\)), which sharpens the model on clean data at the cost of amplifying bad labels.

tsallis_cross_entropy_loss

tsallis_cross_entropy_loss(
    logits_or_probs: Array,
    targets: Array,
    q: Scalar = 1.0,
    from_logits: bool = True,
    normalizer_q: Scalar | None = None,
    axis: int = -1,
    reduction: str = "mean",
) -> Array

q-deformed cross-entropy loss.

Parameters:

Name	Type	Description	Default
logits_or_probs	Array	Unnormalized scores if from_logits (the default), otherwise probabilities that already sum to one along axis.	required
targets	Array	Target distribution along axis. Either one-hot or soft; integer class indices are not accepted, so encode them first.	required
q	Scalar	Entropic index of the loss. 1 is the standard cross-entropy; values below 1 bound the penalty on confidently wrong predictions (robust to label noise); values above 1 sharpen it.	1.0
from_logits	bool	Whether to normalize the input first.	True
normalizer_q	Scalar | None	Entropic index of the tsallis_entmax used to normalize logits. Defaults to q, coupling the two; pass 1.0 to keep an ordinary softmax under a deformed loss. Ignored when from_logits is False.	None
axis	int	Axis holding the class distribution.	-1
reduction	str	"mean", "sum", or "none".	'mean'

Returns:

Type	Description
Array	The reduced loss, or the per-example losses when reduction="none".

Raises:

Type	Description
ValueError	If reduction is not one of the three accepted values.

Note

A sparse normalizer (normalizer_q > 1) can assign exactly zero probability to the true class, for which the loss is a genuine +inf. That is mathematically correct but hostile to optimization; either pair a deformed loss with normalizer_q=1.0 or keep q close to 1 early in training.

Source code in qjax/nn/losses.py

def tsallis_cross_entropy_loss(
    logits_or_probs: Array,
    targets: Array,
    q: Scalar = 1.0,
    from_logits: bool = True,
    normalizer_q: Scalar | None = None,
    axis: int = -1,
    reduction: str = "mean",
) -> jax.Array:
    r"""``q``-deformed cross-entropy loss.

    Args:
        logits_or_probs: Unnormalized scores if ``from_logits`` (the default),
            otherwise probabilities that already sum to one along ``axis``.
        targets: Target distribution along ``axis``. Either one-hot or soft;
            integer class indices are *not* accepted, so encode them first.
        q: Entropic index of the *loss*. ``1`` is the standard cross-entropy;
            values below ``1`` bound the penalty on confidently wrong
            predictions (robust to label noise); values above ``1`` sharpen it.
        from_logits: Whether to normalize the input first.
        normalizer_q: Entropic index of the `tsallis_entmax` used to
            normalize logits. Defaults to ``q``, coupling the two; pass ``1.0``
            to keep an ordinary softmax under a deformed loss. Ignored when
            ``from_logits`` is ``False``.
        axis: Axis holding the class distribution.
        reduction: ``"mean"``, ``"sum"``, or ``"none"``.

    Returns:
        The reduced loss, or the per-example losses when ``reduction="none"``.

    Raises:
        ValueError: If ``reduction`` is not one of the three accepted values.

    Note:
        A sparse normalizer (``normalizer_q > 1``) can assign exactly zero
        probability to the true class, for which the loss is a genuine ``+inf``.
        That is mathematically correct but hostile to optimization; either pair a
        deformed loss with ``normalizer_q=1.0`` or keep ``q`` close to ``1``
        early in training.
    """
    if reduction not in ("mean", "sum", "none"):
        raise ValueError(f"reduction must be 'mean', 'sum', or 'none'; got {reduction!r}.")

    logits_or_probs = jnp.asarray(logits_or_probs, dtype=jnp.result_type(float))
    targets = jnp.asarray(targets, dtype=jnp.result_type(float))

    if from_logits:
        norm_q = q if normalizer_q is None else normalizer_q
        probs = tsallis_entmax(logits_or_probs, q=norm_q, axis=axis)
    else:
        probs = logits_or_probs

    per_example = tsallis_cross_entropy(probs, targets, q, axis=axis)
    if reduction == "mean":
        return jnp.mean(per_example)
    if reduction == "sum":
        return jnp.sum(per_example)
    return per_example

Shared — types, validation, and series

types

Type aliases shared across qjax.

These aliases keep signatures readable without imposing a hard dependency on a specific array backend. At runtime every value is a jax.Array, but the aliases also accept Python scalars and NumPy arrays, which JAX promotes automatically.

Array and Scalar were previously the same alias, so the distinction between "an array of any shape" and "a single real number" was documentation only and a type checker could not act on it. Scalar is now the narrower of the two: it excludes the nested sequences that Array accepts, which is the real constraint on an entropic index.

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).

as_scalar_q

as_scalar_q(q: Scalar) -> Array

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

Parameters:

Name	Type	Description	Default
q	Scalar	The entropic index, as a Python number or array-like.	required

Returns:

Type	Description
Array	A 0-d jax.Array of floating dtype.

Source code in qjax/shared/validation.py

def as_scalar_q(q: Scalar) -> jax.Array:
    """Coerce an entropic index to a floating-point JAX scalar.

    Args:
        q: The entropic index, as a Python number or array-like.

    Returns:
        A 0-d `jax.Array` of floating dtype.
    """
    return jnp.asarray(q, dtype=jnp.result_type(float))

near_one

near_one(q: Scalar, eps: float = Q_EPS) -> Array

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

Parameters:

Name	Type	Description	Default
q	Scalar	The entropic index.	required
eps	float	Half-width of the neighborhood around q = 1.	Q_EPS

Returns:

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

Source code in qjax/shared/validation.py

def near_one(q: Scalar, eps: float = Q_EPS) -> jax.Array:
    """Boolean mask for indices that should use the ``q -> 1`` (classical) limit.

    Args:
        q: The entropic index.
        eps: Half-width of the neighborhood around ``q = 1``.

    Returns:
        A boolean array, broadcastable against ``q``, that is ``True`` where
        ``|q - 1| < eps``.
    """
    return jnp.abs(as_scalar_q(q) - 1.0) < eps

positive_q_or_nan

positive_q_or_nan(q: Scalar) -> Array

Reject a non-positive entropic index.

The Tsallis entropy normalizer \(1/(q(q-1))\) is singular at q = 0, so the entmax family is undefined for q <= 0. A Python-level raise is impossible under jax.jit, so the check is split: a statically known q fails loudly at trace time, while a traced q (e.g. a learnable parameter that wandered out of range) is mapped to NaN so the failure is visible downstream instead of silently plausible.

Parameters:

Name	Type	Description	Default
q	Scalar	The entropic index.	required

Returns:

Type	Description
Array	q as a floating-point JAX scalar, or NaN where a traced q
Array	is non-positive.

Raises:

Type	Description
ValueError	If q is a concrete value and q <= 0.

Source code in qjax/shared/validation.py

def positive_q_or_nan(q: Scalar) -> jax.Array:
    """Reject a non-positive entropic index.

    The Tsallis entropy normalizer $1/(q(q-1))$ is singular at ``q = 0``,
    so the ``entmax`` family is undefined for ``q <= 0``. A Python-level
    ``raise`` is impossible under `jax.jit`, so the check is split: a
    statically known ``q`` fails loudly at trace time, while a traced ``q``
    (e.g. a learnable parameter that wandered out of range) is mapped to ``NaN``
    so the failure is visible downstream instead of silently plausible.

    Args:
        q: The entropic index.

    Returns:
        ``q`` as a floating-point JAX scalar, or ``NaN`` where a traced ``q``
        is non-positive.

    Raises:
        ValueError: If ``q`` is a concrete value and ``q <= 0``.
    """
    q = as_scalar_q(q)
    try:
        static = float(q)
    except (TypeError, jax.errors.ConcretizationTypeError, jax.errors.TracerArrayConversionError):
        return jnp.where(q > 0.0, q, jnp.nan)
    if static <= 0.0:
        raise ValueError(
            f"q must be positive (Tsallis entropy is singular at q = 0); got {static}."
        )
    return q

series

Entire-function ratios used to take the q -> 1 limit stably.

Every Tsallis closed form is a 0/0 indeterminate at q = 1, and the textbook way to evaluate it — select a separate Boltzmann-Gibbs expression when |q - 1| is small — has two defects:

1. Between the switch-over point and the region where the deformed form is accurate there is a band in which neither is: subtracting 1 from x^{1-q} ~ 1 loses most of the mantissa. In float32 the relative error of (x^{1-q} - 1)/(1-q) peaks near 4e-3 around q = 1.00001.
2. The classical branch does not depend on q, so its derivative with respect to q is identically zero. A learnable entropic index that wanders into the switch-over window sees a hard zero gradient.

Both vanish if the indeterminate quotient is written through the entire functions below, which are analytic at the origin and equal 1 there. The classical limit then falls out of the same expression, with no branch on q and no loss of accuracy.

Near the origin each ratio is evaluated by its Taylor series rather than by the direct form, which keeps the value and the derivative correct.

expm1_over_t

expm1_over_t(t: Array) -> Array

Entire function (exp(t) - 1)/t, equal to 1 at t = 0.

Source code in qjax/shared/series.py

def expm1_over_t(t: jax.Array) -> jax.Array:
    """Entire function ``(exp(t) - 1)/t``, equal to ``1`` at ``t = 0``."""
    small = jnp.abs(t) < SERIES_CUTOFF
    safe_t = jnp.where(small, 1.0, t)
    series = 1.0 + t * (0.5 + t * (1.0 / 6.0 + t / 24.0))
    return jnp.where(small, series, jnp.expm1(safe_t) / safe_t)

log1p_over_t

log1p_over_t(t: Array) -> Array

Entire function log(1 + t)/t, equal to 1 at t = 0.

Source code in qjax/shared/series.py

def log1p_over_t(t: jax.Array) -> jax.Array:
    """Entire function ``log(1 + t)/t``, equal to ``1`` at ``t = 0``."""
    small = jnp.abs(t) < SERIES_CUTOFF
    safe_t = jnp.where(small, 1.0, t)
    series = 1.0 - t * (0.5 - t * (1.0 / 3.0 - t * 0.25))
    return jnp.where(small, series, jnp.log1p(safe_t) / safe_t)

Plots

Plotting requires the optional plots extra: pip install "qjax[plots]".

style

Publication-grade plotting style for qjax, themed on the brand ramp.

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

The ramp

QJAX_RAMP is the ten-step green-blue scale the logo and documentation are built from, ordered light to dark. It is registered with Matplotlib as "qjax" (plus "qjax_r"), so CMAP works anywhere a colormap name does.

It is a sequential scale: it encodes magnitude, which is exactly what a family of curves indexed by q needs. Two consequences worth knowing:

• qcolors windows the ramp to [0.40, 1.0] by default. The three lightest steps sit between 1.3:1 and 1.7:1 against white — invisible as thin lines. The window starts where the ramp first clears the 2:1 floor for a sequential light end.
• The ramp cannot supply a categorical palette. Exhaustive search over all 1820 four-colour subsets (including interpolated mid-steps) found none that passes the categorical checks: every subset with usable separation (normal-vision OKLab ΔE >= 15) buys it from the extremes that fall outside the lightness band and below 3:1 contrast. Where a figure distinguishes methods rather than magnitudes, carry identity with linestyle and markers and let colour be the secondary cue.

qcolors

qcolors(
    n: int, lo: float = QCOLORS_LO, hi: float = QCOLORS_HI
) -> list

Sample n evenly spaced colors from the qjax ramp.

Intended for curves indexed by an ordered parameter — a family of q values, a noise sweep — where the reader should see the ordering in the color. For unordered categories (competing methods, class labels) the ramp cannot give reliable separation; see the module docstring.

The default [lo, hi] window starts partway down the ramp so the lightest curve still reads against a white page.

Parameters:

Name	Type	Description	Default
n	int	Number of colors to return.	required
lo	float	Lower bound of the colormap window, in [0, 1].	QCOLORS_LO
hi	float	Upper bound of the colormap window, in [0, 1].	QCOLORS_HI

Returns:

Type	Description
list	A list of n RGBA tuples.

Source code in qjax/plots/style.py

def qcolors(n: int, lo: float = QCOLORS_LO, hi: float = QCOLORS_HI) -> list:
    """Sample ``n`` evenly spaced colors from the qjax ramp.

    Intended for curves indexed by an ordered parameter — a family of ``q``
    values, a noise sweep — where the reader should see the ordering in the
    color. For unordered categories (competing methods, class labels) the ramp
    cannot give reliable separation; see the module docstring.

    The default ``[lo, hi]`` window starts partway down the ramp so the lightest
    curve still reads against a white page.

    Args:
        n: Number of colors to return.
        lo: Lower bound of the colormap window, in ``[0, 1]``.
        hi: Upper bound of the colormap window, in ``[0, 1]``.

    Returns:
        A list of ``n`` RGBA tuples.
    """
    cmap = mpl.colormaps[CMAP]
    if n <= 0:
        return []
    if n == 1:
        # A lone curve should get a mid-ramp color, not the faintest step.
        return [cmap(0.5 * (lo + hi))]
    return [cmap(p) for p in np.linspace(lo, hi, n)]

qlinestyles

qlinestyles(n: int) -> list[LineStyle]

Return n distinguishable Matplotlib linestyles.

The brand ramp is sequential, so colour alone cannot separate more than about three unordered categories: sampling it for four competing methods puts two dark blues side by side that measure well under the readability floor (OKLab ΔE ~6 against a floor of 15). Pairing qcolors with these dash patterns supplies the second, non-colour channel, which also keeps figures readable in grayscale print and for colour-vision deficiencies.

Parameters:

Name	Type	Description	Default
n	int	Number of linestyles to return.	required

Returns:

Type	Description
list[LineStyle]	A list of n linestyle specifications, cycling if n exceeds the
list[LineStyle]	number of defined patterns.

Source code in qjax/plots/style.py

def qlinestyles(n: int) -> list[LineStyle]:
    """Return ``n`` distinguishable Matplotlib linestyles.

    The brand ramp is sequential, so colour alone cannot separate more than about
    three unordered categories: sampling it for four competing *methods* puts two
    dark blues side by side that measure well under the readability floor (OKLab
    ΔE ~6 against a floor of 15). Pairing `qcolors` with these dash patterns
    supplies the second, non-colour channel, which also keeps figures readable in
    grayscale print and for colour-vision deficiencies.

    Args:
        n: Number of linestyles to return.

    Returns:
        A list of ``n`` linestyle specifications, cycling if ``n`` exceeds the
        number of defined patterns.
    """
    if n <= 0:
        return []
    return [QLINESTYLES[i % len(QLINESTYLES)] for i in range(n)]

use_qjax_style

use_qjax_style() -> None

Apply the qjax publication style (serif math, vector PDF, brand-ramp cycle).

Source code in qjax/plots/style.py

def use_qjax_style() -> None:
    """Apply the qjax publication style (serif math, vector PDF, brand-ramp cycle)."""
    plt.rcParams.update(
        {
            # Typography: serif body with Computer-Modern math (no system LaTeX
            # required). pdf/ps fonttype 42 embeds editable TrueType outlines.
            "text.usetex": False,
            "font.family": "serif",
            "font.serif": ["CMU Serif", "Times New Roman", "DejaVu Serif"],
            "mathtext.fontset": "cm",
            "axes.formatter.use_mathtext": True,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
            "font.size": 11,
            "axes.titlesize": 12,
            "axes.labelsize": 11,
            "xtick.labelsize": 9.5,
            "ytick.labelsize": 9.5,
            "legend.fontsize": 9.5,
            # Color: the qjax ramp and a matching discrete cycle.
            "image.cmap": CMAP,
            "axes.prop_cycle": cycler(color=qcolors(5)),
            # Figure / output: single-column default, high-resolution rasters.
            "figure.figsize": (6.0, 4.0),
            "figure.dpi": 150,
            "savefig.dpi": 600,
            "savefig.format": "pdf",
            "savefig.bbox": "tight",
            "savefig.pad_inches": 0.03,
            "savefig.transparent": False,
            # Axes, lines, and ticks.
            "axes.linewidth": 0.8,
            "axes.grid": True,
            "axes.axisbelow": True,
            "grid.alpha": 0.25,
            "grid.linewidth": 0.5,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "lines.linewidth": 1.8,
            "lines.markersize": 5,
            "legend.frameon": False,
            "legend.handlelength": 1.6,
            "xtick.direction": "in",
            "ytick.direction": "in",
            "xtick.minor.visible": True,
            "ytick.minor.visible": True,
            "xtick.major.size": 4.0,
            "ytick.major.size": 4.0,
            "xtick.minor.size": 2.0,
            "ytick.minor.size": 2.0,
        }
    )

save_figure

save_figure(
    fig: Figure, path: str | Path, transparent: bool = False
) -> Path

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:

Name	Type	Description	Default
fig	Figure	The figure to write.	required
path	str | Path	Destination path; any extension is replaced with .pdf.	required
transparent	bool	If True, write with a transparent background so the figure blends with whatever it is placed on.	False

Returns:

Type	Description
Path	The resolved output path.

Source code in qjax/plots/style.py

def save_figure(fig: plt.Figure, path: str | Path, transparent: bool = False) -> Path:
    """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``.

    Args:
        fig: The figure to write.
        path: Destination path; any extension is replaced with ``.pdf``.
        transparent: If ``True``, write with a transparent background so the
            figure blends with whatever it is placed on.

    Returns:
        The resolved output path.
    """
    out = Path(path).with_suffix(".pdf")
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, bbox_inches="tight", pad_inches=0.03, transparent=transparent)
    return out

functions

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

plot_q_log

plot_q_log(
    q_values: Sequence[float] = (0.5, 0.8, 1.0, 1.5, 2.0),
    x_range: tuple[float, float] = (0.05, 4.0),
    num: int = 400,
    ax: Axes | None = None,
) -> Axes

Plot the q-logarithm for several entropic indices.

Parameters:

Name	Type	Description	Default
q_values	Sequence[float]	Entropic indices to draw, one curve each.	(0.5, 0.8, 1.0, 1.5, 2.0)
x_range	tuple[float, float]	(min, max) of the (positive) domain.	(0.05, 4.0)
num	int	Number of sample points.	400
ax	Axes | None	Existing axis to draw on; a new one is created if None.	None

Returns:

Type	Description
Axes	The axis containing the plot.

Source code in qjax/plots/functions.py

def plot_q_log(
    q_values: Sequence[float] = (0.5, 0.8, 1.0, 1.5, 2.0),
    x_range: tuple[float, float] = (0.05, 4.0),
    num: int = 400,
    ax: plt.Axes | None = None,
) -> plt.Axes:
    """Plot the ``q``-logarithm for several entropic indices.

    Args:
        q_values: Entropic indices to draw, one curve each.
        x_range: ``(min, max)`` of the (positive) domain.
        num: Number of sample points.
        ax: Existing axis to draw on; a new one is created if ``None``.

    Returns:
        The axis containing the plot.
    """
    use_qjax_style()
    if ax is None:
        _, ax = plt.subplots()
    x = jnp.linspace(x_range[0], x_range[1], num)
    for q, color in zip(q_values, qcolors(len(q_values)), strict=False):
        ax.plot(x, q_log(x, q), color=color, label=f"q = {q:g}")
    ax.axhline(0.0, color="0.6", lw=0.8)
    ax.set(xlabel="x", ylabel=r"$\ln_q(x)$", title="q-logarithm")
    ax.legend()
    return ax

plot_q_exp

plot_q_exp(
    q_values: Sequence[float] = (0.5, 0.8, 1.0, 1.5, 2.0),
    x_range: tuple[float, float] = (-3.0, 2.0),
    num: int = 400,
    ax: Axes | None = None,
) -> Axes

Plot the q-exponential for several entropic indices.

Parameters:

Name	Type	Description	Default
q_values	Sequence[float]	Entropic indices to draw, one curve each.	(0.5, 0.8, 1.0, 1.5, 2.0)
x_range	tuple[float, float]	(min, max) of the domain.	(-3.0, 2.0)
num	int	Number of sample points.	400
ax	Axes | None	Existing axis to draw on; a new one is created if None.	None

Returns:

Type	Description
Axes	The axis containing the plot.

Source code in qjax/plots/functions.py

def plot_q_exp(
    q_values: Sequence[float] = (0.5, 0.8, 1.0, 1.5, 2.0),
    x_range: tuple[float, float] = (-3.0, 2.0),
    num: int = 400,
    ax: plt.Axes | None = None,
) -> plt.Axes:
    """Plot the ``q``-exponential for several entropic indices.

    Args:
        q_values: Entropic indices to draw, one curve each.
        x_range: ``(min, max)`` of the domain.
        num: Number of sample points.
        ax: Existing axis to draw on; a new one is created if ``None``.

    Returns:
        The axis containing the plot.
    """
    use_qjax_style()
    if ax is None:
        _, ax = plt.subplots()
    x = jnp.linspace(x_range[0], x_range[1], num)
    for q, color in zip(q_values, qcolors(len(q_values)), strict=False):
        ax.plot(x, q_exp(x, q), color=color, label=f"q = {q:g}")
    ax.set(xlabel="x", ylabel=r"$\exp_q(x)$", title="q-exponential")
    ax.legend()
    return ax

distributions

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

plot_q_gaussian

plot_q_gaussian(
    q_values: Sequence[float] = (0.5, 1.0, 1.5, 2.0, 2.5),
    beta: float = 1.0,
    x_range: tuple[float, float] = (-5.0, 5.0),
    num: int = 500,
    ax: Axes | None = None,
) -> Axes

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:

Name	Type	Description	Default
q_values	Sequence[float]	Entropic indices to draw, one curve each (q < 3).	(0.5, 1.0, 1.5, 2.0, 2.5)
beta	float	Shared inverse-width parameter.	1.0
x_range	tuple[float, float]	(min, max) of the domain.	(-5.0, 5.0)
num	int	Number of sample points.	500
ax	Axes | None	Existing axis to draw on; a new one is created if None.	None

Returns:

Type	Description
Axes	The axis containing the plot.

Source code in qjax/plots/distributions.py

def plot_q_gaussian(
    q_values: Sequence[float] = (0.5, 1.0, 1.5, 2.0, 2.5),
    beta: float = 1.0,
    x_range: tuple[float, float] = (-5.0, 5.0),
    num: int = 500,
    ax: plt.Axes | None = None,
) -> plt.Axes:
    """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.

    Args:
        q_values: Entropic indices to draw, one curve each (``q < 3``).
        beta: Shared inverse-width parameter.
        x_range: ``(min, max)`` of the domain.
        num: Number of sample points.
        ax: Existing axis to draw on; a new one is created if ``None``.

    Returns:
        The axis containing the plot.
    """
    use_qjax_style()
    if ax is None:
        _, ax = plt.subplots()
    x = jnp.linspace(x_range[0], x_range[1], num)
    for q, color in zip(q_values, qcolors(len(q_values)), strict=False):
        ax.plot(x, q_gaussian_pdf(x, q, beta), color=color, label=f"q = {q:g}")
    ax.set(xlabel="x", ylabel="density", title=f"q-Gaussian (β = {beta:g})")
    ax.legend()
    return ax