Polynomial Fitting via polyfit

Why this matters

jnp.polyfit(x, y, deg) fits a polynomial of degree deg to (x, y) data via least squares. It returns coefficients in descending power order: c[0] * x^deg + c[1] * x^(deg-1) + ... + c[deg].

This is different from NumPy’s Polynomial class, which uses ascending order. polyfit is used in:

• Trend fitting — fit a line or curve to noisy measurements.
• Smoothing — replace noisy data with a smooth polynomial model.
• Model regression — find polynomial relationships in experimental data.

Key facts:

• For n points and deg = n-1, you get an exact fit (Vandermonde interpolation).
• High-degree fits can be ill-conditioned — scale x to [-1, 1] when needed.
• Output shape is (deg+1,).

Worked mini-example

import jax.numpy as jnp

x = jnp.array([0.0, 1.0, 2.0])
y = jnp.array([0.0, 1.0, 4.0])   # y = x^2
coeffs = jnp.polyfit(x, y, 2)
# coeffs ≈ [1.0, 0.0, 0.0]       # descending: 1*x^2 + 0*x + 0

Common pitfalls

• Ascending vs descending — jnp.polyfit returns coefficients in descending power order; jnp.Polynomial uses ascending. Mixing them silently gives wrong results.
• Ill-conditioning — high-degree polynomials on unscaled x produce huge Vandermonde entries; scale x to [-1, 1] to improve numerics.
• deg must be int — jnp.polyfit requires an integer degree; cast with int(deg) if receiving a float scalar.

Problem

Implement polyfit_coeffs(x, y, deg) that fits a polynomial of degree deg to the data (x, y) via least squares.

• x, y: 1-D jax arrays of the same length.
• deg: scalar (cast to int).
• Returns: 1-D array of shape (deg+1,) — coefficients in descending power order.