layer#
Penzai’s base layer class and associated utilities.
Layer is the base type for most neural network components in Penzai, but also
used for things that are not necessarily neural networks. The key feature of
a Layer is that it can be called with a single positional argument and produces
a single output. (As such, it must implement __call__.)
The single input is allowed to be a collection or data structure, as is the
output. The purpose of the restriction is to make it easier to compose logic
of multiple layers together. In particular, it’s possible to sequentially
run a bunch of layers by passing the output of one layer as the input to the
next. It’s also possible to transform layers by nesting their inputs or outputs
inside structures, without worrying about conflicts. For instance, given a
transformation from a layer x -> y to a layer (x, state) -> (y, state),
and a transformation from a layer x -> y to a layer
(x, random seed) -> y, you can compose these to transform x -> y to
((x, state), random seed) -> (y, state) without the two transformations
having to know about each other.
This overall abstraction is inspired by type polymorphism in functional programming languages, especially as it relates to effect systems.
Layers also make it easier to do shape-checking, by assuming that the input and output are both single PyTrees.
Classes
Abstract base class for neural network layers and other 1-arg callables. |
Functions
|
Decorator for |
|
Shape-checking opt-out decorator for |
Exceptions
Raised when a layer's attributes are misconfigured. |