neural_tangents.stax.serial
- neural_tangents.stax.serial(*layers)[source]
Combinator for composing layers in serial.
Based on jax.example_libraries.stax.serial.
- Parameters
*layers – a sequence of layers, each an (init_fn, apply_fn, kernel_fn) triple.
- Return type
Tuple
[InitFn
,ApplyFn
,LayerKernelFn
]- Returns
A new layer, meaning an (init_fn, apply_fn, kernel_fn) triple, representing the serial composition of the given sequence of layers.