Skip to content

← JAX Numerical Computing step 20 of 25

Easy Primitives

Polynomial Evaluation via polyval

Jump to the task

Turn coefficients into predictions

A polynomial's coefficients tell you how to combine powers of an input. jnp.polyval(coeffs, x_pts) performs that calculation for every point in x_pts. It returns one value per point, with the same shape as x_pts.

Coefficients go from the highest power to the constant term. For example, [2, -3, 1] represents 2*x**2 - 3*x + 1. A zero coefficient still occupies its position: [2, 0, 1] means 2*x**2 + 1, not 2*x + 1. A single coefficient such as [5] represents the constant function 5.

Worked example

import jax.numpy as jnp

coeffs = jnp.array([2.0, -3.0, 1.0])
x_pts = jnp.array([0.0, 1.0, 2.0])
result = jnp.polyval(coeffs, x_pts)
# [1.0, 0.0, 3.0]

At x = 2, the calculation is 2*4 - 3*2 + 1 = 3. At zero, only the constant term remains, so the result is 1.

The output of jnp.polyfit uses this same coefficient order and can be passed directly to jnp.polyval. NumPy's separate numpy.polynomial.Polynomial class uses ascending powers, so its coefficients need a different interpretation. There is no corresponding jnp.Polynomial class to use here.

Problem

Implement polyval_at(coeffs, x_pts) with jnp.polyval.

  • coeffs is a nonempty one-dimensional float32 array, ordered from highest power to constant term.
  • x_pts is a one-dimensional float32 array of evaluation points.
  • Return the polynomial values with the same shape as x_pts. Empty evaluation points give an empty output.

No Python loop or extra vmap is needed. For [1, 0, 0] and points [0, 1, 2, 3], return [0, 1, 4, 9].

Things to watch

Keep zero coefficients for missing powers, and do not reverse the coefficient order. In Python code, ** denotes a power; ^ does not. Array operations return results without changing the supplied arrays. As with other floating-point calculations, high powers or large coefficients can cause rounding error or overflow.