neural_tangents.stax.RectifiedMonomial

neural_tangents.stax.RectifiedMonomial(degree)[source]

Rectified monomials, i.e. (x >= 0) * x^degree.

Parameters:

degree (int) – a non-negative integer power.

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn]

Returns:

(init_fn, apply_fn, kernel_fn).