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, beartype, and jaxtyping (used throughout for runtime shape-checked signatures), 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

Every module's forward is @jaxtyped(typechecker=beartype)-checked for runtime shape/dtype safety, but that decoration is exactly the kind of thing torch.compile can choke on. So each module also takes three constructor-time flags, consistent across all:

• compile: bool = False - if True, the actual math runs through a separate, undecorated _forward_fn wrapped in torch.compile; forward itself stays checked either way, only the inner computation is compiled.
• 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.

PTanhShrink wraps an arbitrary caller-supplied bilinear submodule; if that submodule is itself one of this package's jaxtyped modules, the call into it from inside PTanhShrink's compiled region can still be a graph-break point, fullgraph=True will surface that as a hard error rather than silently falling back.

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.