penzai.nn

Submodules

penzai.nn.attention	Base attention dataflow combinators.
penzai.nn.basic_ops	Basic primitives for neural networks.
penzai.nn.combinators	Dataflow combinators for neural networks.
penzai.nn.dropout	Dropout regularization using RandomEffect.
penzai.nn.embeddings	Structs for handling tokens.
penzai.nn.grouping	Grouping layers, for chaining sequential computations.
penzai.nn.layer	Penzai's base layer and slot classes and associated utilities.
penzai.nn.layer_stack	Layer stacks.
penzai.nn.linear_and_affine	Generalized linear operator layer and associated utilities.
penzai.nn.parameters	Parameters and state variables in Penzai neural networks.
penzai.nn.standardization	Layer normalization (stateless) using named axes.