Metadata-Version: 2.4
Name: stdplus_ext
Version: 0.1.1
Summary: Numerically stable PyTorch stdplus activation for Slashed Normal parameterization
Project-URL: Repository, https://github.com/Yujia-Yan/stdplus_ext
Project-URL: Issues, https://github.com/Yujia-Yan/stdplus_ext/issues
Author: Yujia Yan
License-Expression: MIT
License-File: LICENSE
Keywords: pytorch,slashed-normal,special-functions,stdplus
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.9
Requires-Dist: torch>=2.6
Provides-Extra: dev
Requires-Dist: build>=1; extra == 'dev'
Requires-Dist: mpmath>=1.3; extra == 'dev'
Requires-Dist: pytest>=7; extra == 'dev'
Requires-Dist: twine>=6.1; extra == 'dev'
Description-Content-Type: text/markdown

# 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

```sh
pip install stdplus_ext
```

Requires **Python 3.9+** and **PyTorch 2.6+**. Install your preferred CPU or CUDA
[PyTorch build](https://pytorch.org/get-started/locally/) 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

```python
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:

```python
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](https://openreview.net/pdf?id=6ifeGfWxtX).

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](https://urresearch.rochester.edu/fileDownloadForInstitutionalItem.action?itemFileId=195735&itemId=39584).

## License and source

MIT licensed. [Source code](https://github.com/Yujia-Yan/stdplus_ext) and
[issue tracker](https://github.com/Yujia-Yan/stdplus_ext/issues) are on GitHub.
