# JIT Compilation Optimization Guidelines

## Decision Logic (Framework Router)
Before rewriting, analyze the imports and data structures to select the correct backend:
- IF the code uses `torch.nn` or `torch.Tensor` → Use **PyTorch**.
- IF the code uses `tensorflow`, `keras`, or `tf.Tensor` → Use **TensorFlow**.
- IF the code uses `jax.numpy` or `jax.grad` → Use **JAX**.
- IF the code uses standard `numpy` arrays and Python loops → Use **Numba**.
- IF the code relies heavily on un-compilable Python features (pandas, dynamic dicts, strings) → **Return Original Code**.

## When NOT to Use JIT Compilation

Do NOT generate JIT-compiled candidates in the following scenarios. Instead, focus on other optimization strategies:

### 1. I/O-Bound Functions
Functions that primarily perform file operations, network requests, database queries, or other I/O operations will not benefit from JIT compilation. The bottleneck is external latency, not CPU computation.

### 2. Functions Called Infrequently
JIT compilation has an upfront compilation cost. If a function is only called once or very rarely, the compilation overhead will exceed any runtime savings. JIT is best for hot paths called many times.

### 3. Heavy Use of Unsupported Python Features
Numba's nopython mode does not support:
- Dictionary comprehensions or complex dict operations
- String manipulation (except basic operations)
- List comprehensions with complex logic
- Most Python standard library functions
- Custom Python classes and objects
- Exception handling with complex try/except blocks
- Generators and iterators (limited support)
- Recursion with varying depths (can cause issues)

If the function relies heavily on these features, JIT compilation will either fail or fall back to slow object mode.

### 4. Functions with Dynamic or Heterogeneous Types
JIT compilers optimize for specific type signatures. Functions that receive different types on each call will trigger recompilation, negating performance benefits.

### 5. Small/Trivial Functions
Functions with minimal computation (e.g., simple getters, one-line calculations, basic conditionals) have negligible runtime. JIT overhead is not justified.

### 6. Code Using Incompatible Libraries
Many Python libraries are not JIT-compatible:
- pandas operations (use vectorized pandas methods instead)
- Most third-party libraries
- Complex numpy operations with object dtypes
- Symbolic math libraries (sympy, etc.)

### 7. Functions with Heavy Object Creation
Code that creates many Python objects, uses class instances extensively, or relies on Python's dynamic nature will not JIT well.

## JIT Output Format (CRITICAL)

When you determine that JIT compilation is viable for the given code:
- Apply the JIT decorator **directly** to the output function (e.g., `@numba.njit`, `@torch.compile`, `@tf.function`, `@jax.jit`).
- Add the necessary import (e.g., `import numba`) at the top of the code.
- Do **NOT** create conditional fallback paths. Never wrap JIT usage in `try/except ImportError`, `if HAS_NUMBA`, or any similar if/else branching that falls back to a non-JIT version.
- Do **NOT** create a separate "fast" helper function alongside the original. The output must be a single, clean function with the JIT decorator applied.
- If JIT is not viable for this code, optimize it using other strategies without any JIT decorators — do not include a JIT path at all.

## Guidelines for Numba (`@njit`)
Use Numba when the code:
- Performs numerical computations with NumPy arrays
- Contains loops over array elements
- Uses basic NumPy operations (arithmetic, indexing, slicing)
- Does not use unsupported Python features (classes, dictionaries with non-scalar keys, etc.)

### Numba Best Practices:
- Add `@numba.njit` or `@numba.njit(cache=True)` decorator
- Use `numba.prange` instead of `range` for parallel loops when appropriate
- Avoid Python objects inside JIT functions (use primitive types and arrays)
- Pre-allocate output arrays instead of using list.append()
- Do not use `fastmath=True`
- Ensure all array dtypes are consistent

### Numba Limitations to Avoid:
- No dictionary comprehensions or set operations
- No string operations
- No class instances (use structured arrays instead)
- Limited support for keyword arguments in nested calls

## Guidelines for PyTorch (`torch.compile`)
Use torch.compile when the code:
- Defines or uses PyTorch neural network modules (nn.Module)
- Performs tensor operations with torch tensors
- Contains forward passes or training loops

### torch.compile Best Practices:
- Apply `@torch.compile()` to the model or specific functions
- Use `mode="reduce-overhead"` for small models with many calls
- Use `mode="max-autotune"` for maximum performance when compile time is acceptable
- Consider `fullgraph=True` to ensure the entire function is compiled
- Use `dynamic=True` if input shapes vary

### torch.compile Limitations to Avoid:
- No Python print statements inside compiled functions (causes graph breaks)
- Avoid data-dependent control flow (if statements based on tensor values)
- No unsupported Python built-ins (input(), eval(), exec())
- Avoid in-place operations on input tensors when using certain backends
- No operations with data-dependent output shapes (use `dynamic=True` if needed)
- Avoid calling non-PyTorch functions that can't be traced (use `torch.compiler.disable` to skip them)
- No Python generators or iterators over tensors
- Avoid global variable modifications inside compiled functions
- Limited support for custom autograd functions (may cause graph breaks)
- No third-party libraries that aren't torch-compatible inside the compiled region

## Guidelines for TensorFlow (`tf.function`)
Use tf.function when the code:
- Uses TensorFlow tensors and operations
- Defines or uses Keras models (tf.keras.Model)
- Contains training loops or inference pipelines
- Performs tensor computations with tf.* operations

### tf.function Best Practices:
- Add `@tf.function` decorator to functions performing tensor operations
- Use `jit_compile=True` for XLA compilation: `@tf.function(jit_compile=True)`
- Specify `input_signature` for fixed input shapes to avoid retracing
- Use `experimental_relax_shapes=True` if input shapes vary slightly
- Avoid Python side effects inside tf.function (print, list.append, etc.)
- Use tf.print() instead of print() for debugging inside traced functions

### tf.function Limitations to Avoid:
- No Python control flow that depends on tensor values (use tf.cond, tf.while_loop instead)
- Avoid creating variables inside tf.function (create them outside or in __init__)
- No unsupported Python operations (file I/O, random.random(), etc.)
- Minimize Python object creation inside the function

## Guidelines for JAX (`jax.jit`)
Use jax.jit when the code:
- Uses JAX arrays (jax.numpy operations)
- Performs functional numerical computations
- Contains pure functions without side effects
- Uses JAX transformations (grad, vmap, pmap)

### jax.jit Best Practices:
- Add `@jax.jit` decorator to pure functions
- Use `static_argnums` for arguments that should not be traced (e.g., axis parameters, shapes)
- Use `static_argnames` for keyword arguments that should be static
- Prefer `jax.numpy` over `numpy` for array operations inside jitted functions
- Use `donate_argnums` to allow JAX to reuse input buffers for outputs
- Combine with other JAX transforms: `jax.jit(jax.vmap(fn))` for batched operations

### jax.jit Limitations to Avoid:
- No in-place mutations (JAX arrays are immutable)
- No Python side effects (print, file I/O, global variable modifications)
- No data-dependent control flow with dynamic shapes (use jax.lax.cond, jax.lax.fori_loop instead)
- Avoid Python loops over traced values (use jax.lax.scan or jax.vmap)
- No NumPy arrays inside jitted functions (convert to JAX arrays first)
