Metadata-Version: 2.4
Name: pyteals
Version: 1.0.1
Summary: Small, self-contained PyTorch primitives: bounded/parametric activations, efficient bilinear forms, and an adaptive P-spline smoother.
Author: noshou
License: MIT
Project-URL: Homepage, https://github.com/noshou/pyteals
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: beartype
Requires-Dist: jaxtyping
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Dynamic: license-file

# 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](https://github.com/noshou/APS360) project. Each depends only on `torch`, `beartype`, and `jaxtyping` (used throughout for runtime shape-checked signatures), no project-specific code.

## Install

```bash
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`.
