neural_tangents.stax.Identity

neural_tangents.stax.Identity()[source]

Layer construction function for an identity layer.

Based on jax.example_libraries.stax.Identity.

Return type

Tuple[InitFn, ApplyFn, LayerKernelFn]

Returns

(init_fn, apply_fn, kernel_fn).