layer_stack

Layer stacks.

Classes

LayerStack	A sequence of layers with identical structure, called under jax.lax.scan.
LayerStackGetAttrKey	GetAttrKey for LayerStack with extra metadata.
LayerStackVarBehavior	Behavior of a variable in a layer stack.

Functions

layerstack_axes_from_keypath(keypath)

Extracts the stacked axes from a keypath.