stdplus_ext

Differentiable PyTorch implementations of the stdplus activation function, its logarithm, and their inverses. Slashed Normal uses stdplus to parameterize the standard deviation of a normal posterior distribution through its KL amplitude.

For $x \in \mathbb{R}$, $y = \mathrm{stdplus}(x)$ is the positive solution of

$$ \log(y^2) - y^2 = -x^2 - 1, $$

choosing $y < 1$ for $x < 0$ and $y > 1$ for $x > 0$. It is smooth and increasing, with $\mathrm{stdplus}(0) = 1$.

Installation

pip install stdplus_ext

Requires Python 3.9+ and PyTorch 2.6+. Install your preferred CPU or CUDA PyTorch build first if needed. Supported NVIDIA CUDA devices use Triton acceleration when available. Triton is optional; CPU and installations without Triton use the PyTorch implementation. CPU and NVIDIA CUDA are the tested platforms; other accelerator backends have not been validated.

Quick start

import torch
from stdplus_ext import stdplus, log_stdplus, inv_stdplus, inv_log_stdplus

x = torch.tensor([-2.0, 0.0, 2.0], requires_grad=True)
sigma = stdplus(x)
log_sigma = log_stdplus(x)

recovered = inv_log_stdplus(log_sigma)
sigma.sum().backward()

Functions

• stdplus(x) evaluates the stdplus activation function. Its mathematical range is positive, although sufficiently negative inputs can underflow to zero.
• log_stdplus(x) evaluates $\log\mathrm{stdplus}(x)$ directly. Prefer it for log densities and negative tails to avoid taking the logarithm of an underflowed stdplus result.
• inv_stdplus(y) inverts stdplus for $y > 0$. It returns zero at one, negative infinity at zero, and NaN for negative inputs.
• inv_log_stdplus(z) inverts log_stdplus on the real line and returns zero at zero.

All functions accept dense, strided tensors with float16, bfloat16, float32, or float64 dtype and preserve shape, dtype, and device. Scalars, empty tensors, slices, transposes, and expanded views are supported. Half and bfloat16 inputs use float32 internal arithmetic.

NaNs propagate, and infinities follow the mathematical limits. Results outside the output dtype's range can overflow or underflow. Inverse round trips lose information when the forward value rounds to zero or one; use log space when possible. Python numbers, integer tensors, complex tensors, and sparse tensors are not supported.

Differentiation and compilation

Eager execution supports reverse-mode and forward-mode differentiation, higher derivatives, and torch.func transforms including vmap, jvp, jacrev, and hessian.

Compiled forward and first-order reverse-mode execution are supported:

compiled_stdplus = torch.compile(stdplus, fullgraph=True)

Compilation can improve throughput, especially for large CPU tensors; the first call includes compilation overhead. Combining compilation with higher derivatives or torch.func transforms is not part of the supported contract.

Slashed Normal and KL amplitude

For a standard normal prior $\mathcal{N}(0, 1)$, the Slashed Normal parameterization uses the complex KL amplitude $\psi = a + bi$ to specify the normal posterior $\mathcal{N}(\mu, \sigma^2)$ through

$$ \mu = \sqrt{2},a, \qquad \sigma = \mathrm{stdplus}(\sqrt{2},b) $$

The KL divergence from the posterior to the prior equals the squared modulus of the KL amplitude:

$$ D_{\mathrm{KL}} \left(\mathcal{N}(\mu, \sigma^2),\middle|,\mathcal{N}(0, 1)\right) = |\psi|^2 = a^2 + b^2. $$

References

Yujia Yan, Xingjian Du, and Zhiyao Duan, Slashed Normal: Parameterize Normal Posterior Distributions with KL Amplitude, 2024. Read the paper.

For the updated description, see Yujia Yan, Structured Analysis and Generation in Music, Audio, and Beyond, University of Rochester, 2025, Chapter 6: Section 6.3.2 defines the KL amplitude and stdplus, and Section 6.6 describes the activation function and its numerical evaluation. Read the dissertation.

License and source

MIT licensed. Source code and issue tracker are on GitHub.