layer#
Penzai’s base layer and slot classes and associated utilities.
Layer is the base type for most neural network components in Penzai. Conceptually, a layer is an object that can be called with
a single “ordinary input” (either a single array or a pytree), which usually comes from the previous layer in the model, and is always passed positionally,
and a set of named “side inputs”, which are passed as keyword arguments, represent extra context or state (e.g. attention masks, conditioning information, or random number generators), and are generally forwarded between layers unchanged.
The purpose of this abstraction is to make it easier to compose logic
of multiple layers together. Combinators such as Sequential can be used to
sequence multiple layers, which will feed the output of one layer as the input
to the next, and share the same side inputs across all layers. This simplifies
the “plumbing” of extra information such as attention masks, allowing them to
be used in combination with simple combinators.
Classes
Abstract base class for neural network layers and other 1-arg callables. |