We saw the power of JAX's JIT compiler combined with parallel hardware when we

{ref}`above <jax_speed>`, when we applied `cos` to a large array.

Let's try the same thing with a more complex function.

### Evaluating a more complicated function

Consider the function

Let's try the same thing with a more complex function:

```{code-cell}

def f(x):

y = np.cos(2 * x**2) + np.sqrt(np.abs(x)) + 2 * np.sin(x**4) - x**2

return y

```

#### With NumPy

### With NumPy

We'll try first with NumPy

#### With JAX

### With JAX

Now let's try again with JAX.

second run after JIT compilation.

However, with JAX, we have another trick up our sleeve --- we can JIT-compile

the *entire* function, not just individual operations.

the entire function, not just individual operations.

### Compiling the whole function

### Compiling the Whole Function

The JAX just-in-time (JIT) compiler can accelerate execution within functions by fusing array

operations into a single optimized kernel.