PyTorch Extra Activations And Layers (pyteals)

Small, self-contained PyTorch primitives (currently all nn.Module classes, not guaranteed to stay that way as the package grows), extracted from the ScatterNet project. Each depends only on torch, no project-specific code.

Install

pip install -e /path/to/pyteals   # editable, for local development

Contents

Split into two subpackages by kind: pyteals.activations (pure elementwise nonlinearities, one tensor in, same-shape tensor out) and pyteals.layers (learned/composite layers that take multiple inputs and/or wrap submodules). Everything is also re-exported at the top level, so from pyteals import PBId and from pyteals.activations import PBId both work.

Activations (pyteals.activations)

• PBId - bent identity activation with a learned, bounded per-feature bend strength, interpolating continuously between the identity (w=0) and the standard bent identity (w=1). Use where a layer should be able to learn its way toward linear behaviour rather than being forced through a fixed nonlinearity from the start.
• SqrP - square-plus activation, a smooth strictly-positive approximation to max(x, 0) with a learned width and a numerically stable form on the negative tail. Defined only from +, *, and sqrt (no exp/log), so it's cheaper than softplus on hardware without fast transcendental ops, at the cost of a slower-decaying negative tail (1/|x| vs softplus's exponential); good for a strictly-positive, order-1, softplus-like output where an exact zero is never required. See Barron (2021) below.

Layers (pyteals.layers)

• NoTrilinBilin - drop-in replacement for nn.Bilinear that avoids the generic _trilinear autograd kernel via one matmul plus an elementwise reduce. Best when one side of the bilinear form is small (e.g. min(out_features, in2_features) == 1); profile before reusing it where both sides are large, since it trades PyTorch's fused kernel for an intermediate tensor that PyTorch's implementation never materializes.
• QDiagBilin - bilinear form with one independent weight matrix per point along a designated axis, rather than one matrix shared across it. Use whenever a bilinear combination should vary per grid point, per time step, or per any other structured index instead of being homogeneous across it.
• PTanhShrink - a bilinear layer followed by a width-parametric tanh shrink (y - c*tanh(y/c)), a soft, cubic-near-zero shrink toward 0 with a learned, bounded per-channel width. Generalizes PyTorch's nn.Tanhshrink. Useful wherever a signal should stay near-inert until it clears a per-channel threshold, rather than responding linearly from zero.
• PPSpline - adaptive P-spline / Whittaker smoother for batches of 1D curves, with per-sample, per-point smoothing strength (Λ) supplied at call time (e.g. from an external, input-conditioned "amortized hyperparameter" head) rather than fixed or learned as a free parameter inside the module itself. Use for smoothing a batch of same-length curves where the right amount of smoothing may vary per sample and per position along the curve.

See each module's docstring for the full mathematical description and the reasoning behind its parameterization.

torch.compile

Each module takes three constructor-time flags, consistent across all, controlling whether and how its forward math is compiled:

• compile: bool = False - if True, the actual math runs through a separate _forward_fn wrapped in torch.compile.
• dynamic: bool | None = None - passed straight through to torch.compile. Default None is torch.compile's own default (start static, switch to dynamic shapes automatically on detected recompilation).
• fullgraph: bool = False - passed straight through to torch.compile. Default False falls back to eager on a graph break instead of raising; set True if you'd rather compilation fail loudly on any break.

References

• Barron, J.T. (2021). "Squareplus: A Softplus-Like Algebraic Rectifier." arXiv:2112.11687. https://doi.org/10.48550/arXiv.2112.11687
• PyTorch. torch.nn.Bilinear. https://pytorch.org/docs/stable/generated/torch.nn.Bilinear.html
• PyTorch. torch.nn.Tanhshrink. https://pytorch.org/docs/stable/generated/torch.nn.Tanhshrink.html
• Eilers, P.H.C. & Marx, B.D. (1996). "Flexible smoothing with B-splines and penalties." Statistical Science 11(2), 89-121. DOI 10.1214/ss/1038425655
• Whittaker, E.T. (1922). "On a new method of graduation." Proceedings of the Edinburgh Mathematical Society 41, 63-75.

License

MIT, see LICENSE.