Skip to content
Medium Framework

Compute Gradient

Use automatic differentiation to find the slope of a polynomial at a given point. The function is:

f(x) = x**3 + 2*x**2 + x

Your task

Implement compute_grad(x). The input is a Python number. Return a scalar float32 tensor of shape (), containing the derivative at x. Return the derivative, not the function value or a one-element vector.

For x=1, the function value is 4, but its derivative is 8. The derivative measures how quickly the function changes near that input.

Use the analytical formula 3*x**2 + 4*x + 1 to check your result, but compute the answer with automatic differentiation for this exercise. For x=0 the slope is 1; for x=-1 it is 0.

In PyTorch, create a float32 tensor with requires_grad=True, evaluate the polynomial and call backward() or torch.autograd.grad. Read the derivative from the input, not from the polynomial output.

In JAX, define a function returning the scalar polynomial value and call jax.grad(f) on a float32 input. Keep the computation in tensor operations; converting intermediate values to Python numbers disconnects differentiation.