Topics:
Returns a decorator that augments kernel_fn with consistency checks.
Use this to specify your kernel_fn input kernel requirements.
See also
stax.Diagonal, stax.Input, stax.Output.