fix: enforce direct JIT decorator in optimizer prompt for numerical code

When is_numerical_code is true, the LLM sometimes outputs conditional
fallback paths (try/except, if/else) instead of applying the JIT
decorator directly. Add explicit output format instructions to prevent
this behavior.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
aseembits93 2026-02-23 19:49:24 +05:30
parent 05aecd6fbd
commit 0b523fc367

View file

@ -47,6 +47,15 @@ Many Python libraries are not JIT-compatible:
### 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