Rewrite the following Python code to be JIT-compatible for maximum performance.

Analyze the code to determine which JIT approach to use:
- Numba's `@njit` for NumPy-based numerical code
- `torch.compile` for PyTorch code
- `tf.function` for TensorFlow/Keras code
- `jax.jit` for JAX code

## Source Code

{source_code}

Provide the rewritten JIT-compatible code.
