JAX turns differentiation into a function transformation: jax.grad(f) takes a scalar-valued function and returns a new function that computes $\nabla f$.

Implement df_dx(x) that returns the derivative of

$f(x) = x^3 + 2x$

evaluated at the float x, computed with jax.grad (not by hand-coding $3x^2 + 2$). Return it as a Python float.

Note: jax.grad requires the differentiated input to be a float — passing a Python int raises an error.