neural_tangents.stax.Polynomial

neural_tangents.stax.Polynomial(coef)[source]

Polynomials, i.e. coef[0] + coef[1] * x + + coef[n] * x**n.

Parameters:

coef (Sequence[float]) – a sequence of coefficients. Follows numpy.polynomial.polynomial.Polynomial API.

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn]

Returns:

(init_fn, apply_fn, kernel_fn).