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.



Abstract base class for neural network layers and other 1-arg callables.



Decorator for Layer.__call__ to add shape-checking and name scopes.


Shape-checking opt-out decorator for Layer.__call__.



Raised when a layer's attributes are misconfigured.