neural_tangents.stax.Monomial

neural_tangents.stax.Monomial(degree)[source]

Monomials, i.e. x^degree.

Parameters:

degree (int) – an integer between 0 and 5.

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn]

Returns:

(init_fn, apply_fn, kernel_fn).