stdplus-ext 0.0.1__tar.gz

stdplus_ext-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Yujia Yan
+
+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:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+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.

stdplus_ext-0.0.1/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: stdplus_ext
+Version: 0.0.1
+Summary: The stdplus activation function
+Author: Yujia Yan
+License: MIT
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# stdplus
+
+`stdplus(x)` is a smooth activation function for computing **Gaussian standard deviations** in variational models. It is designed as a numerically stable alternative to `softplus` / `exp`, and is introduced in *Chapter 6* of Yujia Yan’s PhD dissertation (University of Rochester, 2025).  
+It resembles a **1-centered softplus** (i.e., equals 1 at `x = 0`) while providing a clean analytic derivative and a stable evaluation recipe. 
+
+
+## Definition 
+
+stdplus function is defined as a function y = f(x) such that
+$$y^2- 2 \log(y) - 1 = x^2$$
+
+This is proposed in Chapter 6 of Yujia Yan's dissertation: Slashed Normal Parameterization for Approximate Normal Posterior Distributions
+
+
+## Why stdplus?
+
+In the **Slashed Normal / KL-amplitude** parameterization, Gaussian parameters are expressed as a complex “KL amplitude” `ψ = a + bi`, and the standard deviation is computed with:
+
+- `μ = sqrt(2) * a`
+- `σ = stdplus(sqrt(2) * b)` 
+
+This yields an exact closed-form KL term equal to `||ψ||_2^2` (enabling direct, interpretable rate control). 
+
+## Reference
+
+- Yujia Yan, *Structured Analysis and Generation in Music, Audio, and Beyond*, PhD Dissertation, University of Rochester, 2025. (See Chapter 6, Sec. 6.6 “More on stdplus Function”.) 
+
+
+## Install 
+pip3 install -e .
+

stdplus_ext-0.0.1/README.md

@@ -0,0 +1,31 @@
+# stdplus
+
+`stdplus(x)` is a smooth activation function for computing **Gaussian standard deviations** in variational models. It is designed as a numerically stable alternative to `softplus` / `exp`, and is introduced in *Chapter 6* of Yujia Yan’s PhD dissertation (University of Rochester, 2025).  
+It resembles a **1-centered softplus** (i.e., equals 1 at `x = 0`) while providing a clean analytic derivative and a stable evaluation recipe. 
+
+
+## Definition 
+
+stdplus function is defined as a function y = f(x) such that
+$$y^2- 2 \log(y) - 1 = x^2$$
+
+This is proposed in Chapter 6 of Yujia Yan's dissertation: Slashed Normal Parameterization for Approximate Normal Posterior Distributions
+
+
+## Why stdplus?
+
+In the **Slashed Normal / KL-amplitude** parameterization, Gaussian parameters are expressed as a complex “KL amplitude” `ψ = a + bi`, and the standard deviation is computed with:
+
+- `μ = sqrt(2) * a`
+- `σ = stdplus(sqrt(2) * b)` 
+
+This yields an exact closed-form KL term equal to `||ψ||_2^2` (enabling direct, interpretable rate control). 
+
+## Reference
+
+- Yujia Yan, *Structured Analysis and Generation in Music, Audio, and Beyond*, PhD Dissertation, University of Rochester, 2025. (See Chapter 6, Sec. 6.6 “More on stdplus Function”.) 
+
+
+## Install 
+pip3 install -e .
+

stdplus_ext-0.0.1/pyproject.toml

@@ -0,0 +1,20 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "stdplus_ext"                      
+version = "0.0.1"
+description = "The stdplus activation function"
+readme = "README.md"
+requires-python = ">=3.6"
+license = {text = "MIT"}
+authors = [{name = "Yujia Yan"}]
+
+dependencies = []
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

stdplus_ext-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

stdplus_ext-0.0.1/src/stdplus_ext/__init__.py

@@ -0,0 +1,4 @@
+from .stdplus import stdplus, log_stdplus, inv_log_stdplus, inv_stdplus
+
+
+__all__ = ["stdplus", "log_stdplus", "inv_stdplus", "inv_log_stdplus"]

stdplus_ext-0.0.1/src/stdplus_ext/log_stdplus_triton.py

@@ -0,0 +1,173 @@
+import math
+
+import torch
+from torch.utils._triton import has_triton
+
+import triton
+from triton import language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=32, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=32, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=16, num_stages=1),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=32, num_stages=1),
+    ],
+    key=[],
+)
+@triton.jit
+def log_stdplus_kernel(
+    in_ptr0,
+    out_ptr,
+    out_d_ptr,
+    n_elements,
+    n_iterations,
+    BLOCK_SIZE: "tl.constexpr",
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < n_elements
+
+    x = tl.load(in_ptr0 + offsets, mask=mask)
+
+    small_mask = (x < 0.04) & (x > -0.04)
+
+    b1 = 0.7071067811865476  # M_SQRT1_2
+    b2 = 0.25  # 1.0/4
+    b3 = 0.007856742013183984  # M_SQRT1_2/90.0
+    a1 = 0.5892556509887896  # 5.0*M_SQRT1_2/6.0
+    a2 = 0.09444444444444444  # 17.0/180
+
+    x_small = tl.where(small_mask, x, 0.0)
+    fx = x_small * (b1 + x_small * (b2 + x_small * b3))
+    d_fx = b1 + x_small * (2.0 * b2 + 3.0 * x_small * b3)
+    gx = 1.0 + x_small * (a1 + x_small * (a2))
+    d_gx = a1 + 2.0 * x_small * (a2)
+
+    result = fx / gx
+
+    result_d = (d_fx * gx - fx * d_gx) / (gx * gx)
+
+    # branch2
+    regular_mask = ~small_mask
+    x = tl.where(regular_mask, x, 1.0)
+
+    xSqr_p_1 = x * x + 1
+    logy = tl.math.log((x + tl.math.sqrt(x * x + 4.0)) / 2.0) * 2.0
+
+    is_positive = x >= 0
+
+    for i in range(n_iterations):
+        exp_nabslogy = tl.exp(-tl.abs(logy))
+
+        exp_r_a = tl.where(is_positive, 1.0, exp_nabslogy)
+        exp_a = tl.where(is_positive, exp_nabslogy, 1.0)
+
+        logy = ((logy - 1) * exp_r_a + xSqr_p_1 * exp_a) / (exp_r_a - exp_a)
+
+    d_logy = x / (tl.math.exp(logy) - 1)
+    logy = logy / 2
+
+    result = tl.where(regular_mask, logy, result)
+    result_d = tl.where(regular_mask, d_logy, result_d)
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+    tl.store(out_d_ptr + offsets, result_d, mask=mask)
+
+
+def _log_stdplus_triton(x, eps: float):
+
+    n_iterations = int(
+        math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1)
+    )
+
+    output = torch.empty_like(x)
+    output_d = torch.empty_like(x)
+
+    n_elements = output.numel()
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    log_stdplus_kernel[grid](x, output, output_d, n_elements, n_iterations)
+
+    return output, output_d
+
+
+class _LOG_STDPLUS_TRITON(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x):
+        eps = torch.finfo(x.dtype).eps
+
+        logy, d_logy = _log_stdplus_triton(x, eps)
+
+        ctx.save_for_backward(d_logy)
+
+        return logy
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (derivative,) = ctx.saved_tensors
+
+        return grad_output * derivative
+
+
+log_stdplus_triton = _LOG_STDPLUS_TRITON.apply
+
+
+if __name__ == "__main__":
+
+    import matplotlib.pyplot as plt
+    import time
+    import torch.nn.functional as F
+    import pytorch_stdplus
+
+    # x = torch.linspace(-0.08,0.08,65535,  device="cuda")
+    # x = torch.linspace(-1.00,1.00,65535,  device="cuda")
+    x = torch.linspace(-5.00, 5.00, 65535, device="cuda")
+
+    x.requires_grad_()
+
+    out = log_stdplus_triton(x).exp()
+
+    t1 = time.time()
+    for i in range(1000):
+        out = F.softplus(x)
+        # out = torch.log(torch.exp(x)+1)
+    t2 = time.time()
+    print(t2 - t1)
+
+    t1 = time.time()
+    for i in range(1000):
+        yGT = log_stdplus_triton(x)  # .exp()
+    t2 = time.time()
+    print(t2 - t1)
+
+    t1 = time.time()
+    for i in range(1000):
+        out = pytorch_stdplus.log_stdplus(x)  # .exp()
+    t2 = time.time()
+    print(t2 - t1)
+
+    # y**2 - 2*ln(y) - 1 = x**2
+    # torch.expm(out*2)  - 2*out - 1
+
+    # import slashed_normal
+    # yGT  = slashed_normal.log_stdplus(x)
+
+    d_GT = torch.autograd.grad(yGT.sum(), x)[0]
+    out_d = torch.autograd.grad(out.sum(), x)[0]
+
+    plt.plot(x.cpu().detach(), out.cpu().detach())
+    plt.plot(x.cpu().detach(), out_d.cpu().detach())
+    plt.show()
+    print((yGT - out).abs().max())
+    print((d_GT - out_d).abs().max())

stdplus_ext-0.0.1/src/stdplus_ext/stdplus.py

@@ -0,0 +1,246 @@
+import torch
+import math
+
+
+def handles_scalar_tensor(func):
+    def wrapper(x):
+        if torch.is_tensor(x) and x.dim() == 0:
+            x = x.unsqueeze(-1)
+            return func(x).squeeze(-1)
+        else:
+            return func(x)
+
+    return wrapper
+
+
+@torch.jit.script
+def _log_stdplus(x, eps: float, exponential: bool = False):
+
+    # iters = int(math.ceil(math.log(- math.log(eps)/math.log(2))/math.log(2)-1))
+    iters = 4
+
+    # Pre-compute constants once
+    b1 = 1.0 / math.sqrt(2)
+    b2 = 0.25
+    b3 = b1 / 90.0
+    a1 = 5.0 * b1 / 6.0
+    a2 = 17.0 / 180.0
+
+    # Compute all intermediate values in one go
+    x_sq = x * x
+
+    # Small values optimization (-0.04 < x < 0.04)
+    small_vals_mask = (x > -0.04) & (x < 0.04)
+
+    # Initialize result tensor
+    result = torch.empty_like(x)
+
+    # Handle small values first
+    if small_vals_mask.any():
+        # Horner's method for polynomial evaluation
+        fx = x[small_vals_mask] * (
+            b1 + x[small_vals_mask] * (b2 + x[small_vals_mask] * b3)
+        )
+        gx = 1.0 + x[small_vals_mask] * (a1 + x[small_vals_mask] * a2)
+        small_result = fx / gx
+
+        result[small_vals_mask] = small_result
+
+    # Handle non-small values
+    non_small_mask = ~small_vals_mask
+    if non_small_mask.any():
+        # Compute initial values for non-small x
+        x_non_small = x[non_small_mask]
+        x_sq_non_small = x_sq[non_small_mask]
+
+        # Optimize sqrt computation
+        # xsqr_p_4 = torch.sqrt(x_sq_non_small + 4.0)
+        xsqr_p_4 = torch.hypot(x_non_small, torch.ones_like(x_non_small) * 2.0)
+
+        tmp = torch.where(
+            x_non_small < 0, 4.0 / (xsqr_p_4 - x_non_small), x_non_small + xsqr_p_4
+        )
+
+        # logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
+        logy = 2.0 * torch.log(tmp * 0.5)
+
+        xSqr_p_1 = x_sq_non_small + 1.0
+        # print(logy)
+
+        # Optimization: pre-compute is_positive
+        is_positive = x_non_small >= 0
+
+        # Main iteration loop
+        for _ in range(iters):
+            abs_logy = torch.abs(logy)
+            exp_nabslogy = torch.exp(-abs_logy)
+
+            # Vectorized condition handling
+            exp_r_a = torch.where(is_positive, torch.ones_like(logy), exp_nabslogy)
+            exp_a = torch.where(is_positive, exp_nabslogy, torch.ones_like(logy))
+
+            # Combined computation
+            numerator = (logy - 1) * exp_r_a + xSqr_p_1 * exp_a
+            denominator = exp_r_a - exp_a
+
+            # Update in-place
+            logy = numerator / denominator
+
+        # Final computations
+        logy = logy * 0.5
+
+        result[non_small_mask] = logy
+
+    if exponential:
+        result = torch.exp(result)
+
+    return result
+
+
+@torch.jit.script
+def _d_log_stdplus(x, eps: float, exponential: bool = False):
+    """
+    Optimized PyTorch implementation of log_stdplus function.
+
+    Args:
+        x (torch.Tensor): Input tensor
+        iters (int): Number of iterations for the approximation
+        exponential (bool): If True, returns exp(result) for stdplus computation
+
+    Returns:
+        torch.Tensor: Computed result
+    """
+    iters = int(math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1))
+
+    # Pre-compute constants once
+    b1 = 1.0 / math.sqrt(2)
+    b2 = 0.25
+    b3 = b1 / 90.0
+    a1 = 5.0 * b1 / 6.0
+    a2 = 17.0 / 180.0
+
+    # Compute all intermediate values in one go
+    x_sq = x * x
+
+    # Small values optimization (-0.04 < x < 0.04)
+    small_vals_mask = (x > -0.04) & (x < 0.04)
+
+    # Initialize result tensor
+    result = torch.empty_like(x)
+
+    # Handle small values first
+    if small_vals_mask.any():
+        # Horner's method for polynomial evaluation
+        fx = x[small_vals_mask] * (
+            b1 + x[small_vals_mask] * (b2 + x[small_vals_mask] * b3)
+        )
+        gx = 1.0 + x[small_vals_mask] * (a1 + x[small_vals_mask] * a2)
+        d_fx = b1 + x[small_vals_mask] * (2.0 * b2 + x[small_vals_mask] * (3.0 * b3))
+        d_gx = a1 + x[small_vals_mask] * (2.0 * a2)
+
+        small_result = (d_fx * gx - fx * d_gx) / gx / gx
+
+        if exponential:
+            small_result = torch.exp(fx / gx) * small_result
+
+        result[small_vals_mask] = small_result
+
+    # Handle non-small values
+    non_small_mask = ~small_vals_mask
+    if non_small_mask.any():
+        # Compute initial values for non-small x
+        x_non_small = x[non_small_mask]
+        x_sq_non_small = x_sq[non_small_mask]
+
+        # Optimize sqrt computation
+        # xsqr_p_4 = torch.sqrt(x_sq_non_small + 4.0)
+        # logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
+        xsqr_p_4 = torch.hypot(x_non_small, torch.ones_like(x_non_small) * 2.0)
+        tmp = torch.where(
+            x_non_small < 0, 4.0 / (xsqr_p_4 - x_non_small), x_non_small + xsqr_p_4
+        )
+
+        # logy = 2.0 * torch.log((x_non_small + xsqr_p_4) * 0.5)
+        logy = 2.0 * torch.log(tmp * 0.5)
+        xSqr_p_1 = x_sq_non_small + 1.0
+
+        # Optimization: pre-compute is_positive
+        is_positive = x_non_small >= 0
+
+        # Main iteration loop
+        for _ in range(iters):
+            abs_logy = torch.abs(logy)
+            exp_nabslogy = torch.exp(-abs_logy)
+
+            # Vectorized condition handling
+            exp_r_a = torch.where(is_positive, torch.ones_like(logy), exp_nabslogy)
+            exp_a = torch.where(is_positive, exp_nabslogy, torch.ones_like(logy))
+
+            # Combined computation
+            numerator = (logy - 1) * exp_r_a + xSqr_p_1 * exp_a
+            denominator = exp_r_a - exp_a
+
+            # Update in-place
+            logy = numerator / denominator
+
+        # Final computations
+        logy = logy * 0.5
+
+        d_logy = x_non_small / torch.expm1(logy * 2)
+        if exponential:
+            d_logy = torch.exp(logy) * d_logy
+
+        result[non_small_mask] = d_logy
+
+    return result
+
+
+from torch.utils._triton import has_triton as _has_triton
+
+if _has_triton:
+    from .log_stdplus_triton import log_stdplus_triton as log_stdplus_triton
+    from .stdplus_triton import stdplus_triton as stdplus_triton
+
+
+class _LOG_STDPLUS(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        eps = torch.finfo(x.dtype).eps
+        return _log_stdplus(x, eps, exponential=False)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        eps = torch.finfo(x.dtype).eps
+        derivative = _d_log_stdplus(x, eps, exponential=False)
+
+        return grad_output * derivative
+
+
+@handles_scalar_tensor
+def stdplus(x):
+    if x.is_cuda and _has_triton:
+        return stdplus_triton(x)
+
+    return _LOG_STDPLUS.apply(x).exp()
+
+
+@handles_scalar_tensor
+def log_stdplus(x):
+    if x.is_cuda and _has_triton:
+        return log_stdplus_triton(x)
+    return _LOG_STDPLUS.apply(x)
+
+
+@handles_scalar_tensor
+def inv_log_stdplus(y):
+
+    y = y * 2
+    return torch.sgn(y) * (torch.expm1(y) - y).sqrt()
+
+
+@handles_scalar_tensor
+def inv_stdplus(y):
+    return inv_log_stdplus(y.log())

stdplus_ext-0.0.1/src/stdplus_ext/stdplus_triton.py

@@ -0,0 +1,205 @@
+import math
+
+import torch
+from torch.utils._triton import has_triton
+
+import triton
+from triton import language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 64}, num_warps=32, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 128}, num_warps=32, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=8, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=16, num_stages=4),
+        triton.Config(kwargs={"BLOCK_SIZE": 256}, num_warps=32, num_stages=4),
+    ],
+    key=["n_iterations"],
+)
+@triton.jit
+def stdplus_kernel(
+    in_ptr0,
+    out_ptr,
+    out_d_ptr,
+    n_elements,
+    n_iterations,
+    BLOCK_SIZE: "tl.constexpr",
+):
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < n_elements
+
+    x = tl.load(in_ptr0 + offsets, mask=mask)
+
+    small_mask = (x < 0.04) & (x > -0.04)
+
+    # firstly computes log_stdplus
+
+    # b1 = 1.0/tl.math.sqrt(2.0)
+    # b2 = 1.0/4
+    # b3 = 1.0/(90*tl.math.sqrt(2.0))
+
+    # a1 = 5.0/(6*tl.math.sqrt(2.0))
+
+    b1 = 0.7071067811865476  # M_SQRT1_2
+    b2 = 0.25  # 1.0/4
+    b3 = 0.007856742013183984  # M_SQRT1_2/90.0
+    a1 = 0.5892556509887896  # 5.0*M_SQRT1_2/6.0
+    a2 = 0.09444444444444444  # 17.0/180
+
+    # result = tl.zeros_like(x)
+    # result_d = tl.zeros_like(x)
+
+    # logy = x1*(b1+ x1*(b2+ x1*b3))  \
+    # /(1+x1*(a1+ x1*(a2)))
+    x_small = tl.where(small_mask, x, 0.0)
+    fx = x_small * (b1 + x_small * (b2 + x_small * b3))
+    d_fx = b1 + x_small * (2.0 * b2 + 3.0 * x_small * b3)
+    gx = 1.0 + x_small * (a1 + x_small * (a2))
+    d_gx = a1 + 2.0 * x_small * (a2)
+
+    # result = tl.where(small_mask, fx/gx, result)
+
+    # result_d = tl.where(small_mask, (d_fx*gx - fx*d_gx)/ (gx*gx), result_d)
+    result = fx / gx
+
+    result_d = (d_fx * gx - fx * d_gx) / (gx * gx)
+
+    # branch2
+    regular_mask = ~small_mask
+    x = tl.where(regular_mask, x, 1.0)
+
+    xSqr_p_1 = x * x + 1
+    logy = tl.math.log((x + tl.math.sqrt(x * x + 4.0)) / 2.0) * 2.0
+
+    # for i in range(n_iterations):
+
+    # a = tl.math.max(logy, 0.0)
+    # exp_r_a = tl.math.exp(logy-a)
+    # exp_a = tl.math.exp(-a)
+    # logy = ((logy-1)*exp_r_a + xSqr_p_1* exp_a)/( exp_r_a - exp_a)
+
+    is_positive = x >= 0
+
+    # for i in range(n_iterations):
+    for i in range(n_iterations):
+        exp_nabslogy = tl.exp(-tl.abs(logy))
+
+        exp_r_a = tl.where(is_positive, 1.0, exp_nabslogy)
+        exp_a = tl.where(is_positive, exp_nabslogy, 1.0)
+
+        logy = ((logy - 1) * exp_r_a + xSqr_p_1 * exp_a) / (exp_r_a - exp_a)
+
+    logy = logy / 2
+    # expm1_2_logy = torch.expm1(logy*2)
+    d_logy = x / (tl.math.exp(logy * 2) - 1)
+
+    result = tl.where(regular_mask, logy, result)
+    result_d = tl.where(regular_mask, d_logy, result_d)
+
+    ## transform
+    result = tl.math.exp(result)
+    result_d = result_d * result
+
+    tl.store(out_ptr + offsets, result, mask=mask)
+    tl.store(out_d_ptr + offsets, result_d, mask=mask)
+
+
+def _stdplus_triton(x, eps: float):
+
+    n_iterations = int(
+        math.ceil(math.log(-math.log(eps) / math.log(2)) / math.log(2) - 1)
+    )
+
+    output = torch.empty_like(x)
+    output_d = torch.empty_like(x)
+
+    n_elements = output.numel()
+
+    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    stdplus_kernel[grid](x, output, output_d, n_elements, n_iterations)
+
+    return output, output_d
+
+
+class _STDPLUS_TRITON(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x):
+        eps = torch.finfo(x.dtype).eps
+
+        logy, d_logy = _stdplus_triton(x, eps)
+
+        ctx.save_for_backward(d_logy)
+
+        return logy
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (derivative,) = ctx.saved_tensors
+
+        return grad_output * derivative
+
+
+stdplus_triton = _STDPLUS_TRITON.apply
+
+
+if __name__ == "__main__":
+
+    import matplotlib.pyplot as plt
+    import time
+    import torch.nn.functional as F
+    import pytorch_stdplus
+
+    # x = torch.linspace(-0.08,0.08,65535,  device="cuda")
+    # x = torch.linspace(-1.00,1.00,65535,  device="cuda")
+    x = torch.linspace(-5.00, 5.00, 6553500, device="cuda")
+
+    x.requires_grad_()
+
+    out = stdplus_triton(x)
+
+    print("gelu")
+    t1 = time.time()
+    for i in range(10000):
+        out1 = F.gelu(x)
+        # out = torch.log(torch.exp(x)+1)
+    t2 = time.time()
+    print("gelu", t2 - t1)
+
+    t1 = time.time()
+    for i in range(10000):
+        out2 = stdplus_triton(x)  # .exp()
+    t2 = time.time()
+    print("triton", t2 - t1)
+
+    out3 = pytorch_stdplus.stdplus(x)  # .exp()
+    out3 = pytorch_stdplus.stdplus(x)  # .exp()
+    t1 = time.time()
+    for i in range(10000):
+        out3 = pytorch_stdplus.stdplus(x)  # .exp()
+    t2 = time.time()
+    print("cuda", t2 - t1)
+
+    print((out3**2 - 2 * torch.log(out3) - 1 - x**2).abs().max())
+
+    print((out3 - out2).abs().max())
+    # import slashed_normal
+    # yGT  = slashed_normal.log_stdplus(x)
+
+    out_d_2 = torch.autograd.grad(out2.sum(), x)[0]
+    out_d_3 = torch.autograd.grad(out3.sum(), x)[0]
+    print((out_d_2 - out_d_3).abs().max())
+
+    plt.plot(x.cpu().detach(), out2.cpu().detach())
+    plt.plot(x.cpu().detach(), out_d_2.cpu().detach())
+    plt.show()
+    # print((yGT-out).abs().max())

stdplus_ext-0.0.1/src/stdplus_ext.egg-info/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: stdplus_ext
+Version: 0.0.1
+Summary: The stdplus activation function
+Author: Yujia Yan
+License: MIT
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# stdplus
+
+`stdplus(x)` is a smooth activation function for computing **Gaussian standard deviations** in variational models. It is designed as a numerically stable alternative to `softplus` / `exp`, and is introduced in *Chapter 6* of Yujia Yan’s PhD dissertation (University of Rochester, 2025).  
+It resembles a **1-centered softplus** (i.e., equals 1 at `x = 0`) while providing a clean analytic derivative and a stable evaluation recipe. 
+
+
+## Definition 
+
+stdplus function is defined as a function y = f(x) such that
+$$y^2- 2 \log(y) - 1 = x^2$$
+
+This is proposed in Chapter 6 of Yujia Yan's dissertation: Slashed Normal Parameterization for Approximate Normal Posterior Distributions
+
+
+## Why stdplus?
+
+In the **Slashed Normal / KL-amplitude** parameterization, Gaussian parameters are expressed as a complex “KL amplitude” `ψ = a + bi`, and the standard deviation is computed with:
+
+- `μ = sqrt(2) * a`
+- `σ = stdplus(sqrt(2) * b)` 
+
+This yields an exact closed-form KL term equal to `||ψ||_2^2` (enabling direct, interpretable rate control). 
+
+## Reference
+
+- Yujia Yan, *Structured Analysis and Generation in Music, Audio, and Beyond*, PhD Dissertation, University of Rochester, 2025. (See Chapter 6, Sec. 6.6 “More on stdplus Function”.) 
+
+
+## Install 
+pip3 install -e .
+

stdplus_ext-0.0.1/src/stdplus_ext.egg-info/SOURCES.txt

@@ -0,0 +1,11 @@
+LICENSE
+README.md
+pyproject.toml
+src/stdplus_ext/__init__.py
+src/stdplus_ext/log_stdplus_triton.py
+src/stdplus_ext/stdplus.py
+src/stdplus_ext/stdplus_triton.py
+src/stdplus_ext.egg-info/PKG-INFO
+src/stdplus_ext.egg-info/SOURCES.txt
+src/stdplus_ext.egg-info/dependency_links.txt
+src/stdplus_ext.egg-info/top_level.txt

stdplus_ext-0.0.1/src/stdplus_ext.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

stdplus_ext-0.0.1/src/stdplus_ext.egg-info/top_level.txt

@@ -0,0 +1 @@
+stdplus_ext