# JIT Compilation Test Writing Guidelines

When generating tests for functions that may be optimized with JIT compilation (numba `@njit`, `@torch.compile`, `@tf.function`, `@jax.jit`), follow these guidelines to ensure tests remain valid after optimization:

## Critical: Avoid Testing Error Behavior

**DO NOT write tests that check for specific exception types or error handling behavior.** JIT compilers produce different error types than the uncompiled versions:

- Numba raises `TypingError` instead of `TypeError` for invalid inputs
- `torch.compile` may auto-promote types instead of raising errors
- Empty tensors produce warnings instead of `RuntimeError`
- Mismatched array lengths may raise different exception types

**Instead of testing error cases, focus exclusively on testing valid inputs and correct outputs.**

### What NOT to Test

```python
# BAD: Testing for specific exceptions - these differ between compiled/uncompiled
def test_none_raises_type_error():
    with pytest.raises(TypeError):
        func(None)  # Numba raises TypingError, not TypeError

def test_empty_raises_runtime_error():
    with pytest.raises(RuntimeError):
        func(torch.tensor([]))  # May return NaN with warning instead

def test_list_raises_type_error():
    with pytest.raises(TypeError):
        func([1, 2, 3])  # Numpy actually accepts lists

def test_mismatched_shapes_raise():
    with pytest.raises(ValueError):
        func(arr1, arr2)  # Different exception types in JIT
```

### What TO Test

```python
# GOOD: Test valid inputs and their expected outputs
def test_basic_computation():
    arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
    result = func(arr)
    assert np.allclose(result, expected, rtol=1e-7)

def test_edge_case_empty_returns_correct_shape():
    arr = np.empty((0, 3), dtype=np.float64)
    result = func(arr)
    assert result.shape == (0,)  # Test behavior, not exceptions

def test_single_element():
    arr = np.array([5.0], dtype=np.float64)
    result = func(arr)
    assert np.allclose(result, 5.0)
```

## Key Principles

### 1. Use Concrete, Typed Inputs — With Consistent Shapes and Dtypes
JIT compilers trace and compile a new graph for each unique combination of tensor rank, shape, and dtype. Every variation triggers a **recompilation**, which is expensive and defeats the purpose of JIT. To minimize recompilations:
- Use numpy arrays instead of Python lists for numerical data
- Specify dtypes explicitly (e.g., `np.array([1, 2, 3], dtype=np.float64)`)
- **Pick ONE canonical shape and dtype and reuse them across all test cases for the same parameter.** Only vary shapes/dtypes in a dedicated test if that variation is the point of the test
- Avoid mixing types in the same parameter across different test cases
- **Never use scalar tensors (0-d) in some tests and 1-d/2-d tensors in others for the same parameter** — this changes the tensor rank and always triggers recompilation

### 2. Avoid Python Object Features
JIT-compiled functions have limited support for Python objects. In tests:
- Do not test with custom Python class instances as inputs (unless the function is clearly designed for them)
- Avoid dictionary inputs with dynamic keys
- Prefer numpy arrays and primitive types (int, float, bool)

### 3. Handle Numerical Precision
JIT compilation may produce slightly different floating-point results. When testing numerical functions:
- Use `np.allclose()`/`torch.allclose()`/`jax.numpy.allclose()`/`tf.experimental.numpy.allclose(a, b)`/`pytest.approx()` for floating-point comparisons instead of exact equality
- Allow for small numerical tolerances (e.g., `rtol=1e-7, atol=1e-10`)

### 4. Do Not Mock JIT-Decorated Functions
Never mock or patch JIT-decorated functions directly. Instead:
- Test the function's actual behavior
- Mock only external I/O dependencies if absolutely necessary

### 5. Large Scale Tests Should Use Appropriate Data Types
For performance/scalability tests with large data:
- Use contiguous arrays if possible.
- Ensure array dtypes are JIT-compatible (float64, float32, int64, int32, etc.)
- Avoid object dtypes in arrays

## Recompilation Prevention (CRITICAL for torch.compile / tf.function / jax.jit)

JIT compilers trace a new graph for each unique combination of tensor rank, shape, and dtype. Every variation triggers an expensive **recompilation**. You MUST follow these rules:

1. **Pick one canonical shape per parameter and reuse it in every test.** For example, if the function takes a 1-D tensor, use `(4,)` everywhere. If it takes a 2-D tensor, use `(4, 3)` everywhere. Do NOT vary shapes across test functions.
2. **Use `torch.float32` as the default dtype for all tests.** Only use a different dtype if the specific test is *about* dtype behavior. Never use `int32`, `int64`, or `float64` unless required by the function signature.
3. **Never mix tensor ranks across tests.** If one test passes a 1-D tensor `(5,)`, do not pass a scalar tensor `()`, a 2-D tensor `(3, 4)`, or a 0-element tensor `(0,)` in other tests — each triggers a recompilation.
4. **Keep edge-case shapes in a single dedicated test.** If you need to test empty tensors, large tensors, or different ranks, group them in one test function (e.g., `test_edge_cases`) so the recompilation cost is isolated.
5. **Do not pass integer-dtype tensors** (`int32`, `int64`) to functions that use `mean()`, `var()`, or other float-only ops — `torch.compile` will fail with a dtype error.
6. **Do not pass `None` as a tensor argument** — `torch.compile` will trace it and crash.

## Example Test Patterns

### Good: Typed numpy input
```python
def test_compute_sum():
    arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
    result = compute_sum(arr)
    assert np.isclose(result, 6.0)
```

### Good: Tolerance-based comparison
```python
def test_matrix_operation():
    a = np.random.rand(100, 100).astype(np.float64)
    b = np.random.rand(100, 100).astype(np.float64)
    result = matrix_multiply(a, b)
    expected = np.dot(a, b)
    assert np.allclose(result, expected, rtol=1e-7)
```

### Avoid: Python list input for numerical functions
```python
# Avoid this pattern for functions that may be JIT-compiled
def test_compute_sum():
    result = compute_sum([1, 2, 3])  # Python list - may cause JIT issues
```

### Avoid: Mixed types or shapes across tests
```python
# Avoid mixing types across test cases
def test_function():
    result1 = func(np.array([1, 2, 3]))  # int array
    result2 = func(np.array([1.0, 2.0]))  # float array - different type signature

# Avoid: different shapes/ranks in each test — each triggers a recompilation
def test_a():
    func(np.array([1.0, 2.0, 3.0]))           # shape (3,) - 1st trace
def test_b():
    func(np.array([[1.0, 2.0], [3.0, 4.0]]))  # shape (2,2) - recompilation!
```

---

## Framework-Specific Guidelines

### Numba (`@njit`, `@jit`)

**Supported Types:**
- Primitive types: `int`, `float`, `bool`, `complex`
- NumPy arrays with numeric dtypes (`float32`, `float64`, `int32`, `int64`, `complex64`, `complex128`)
- NumPy scalars
- Tuples of supported types (homogeneous preferred)
- Named tuples (with type annotations)

**Test Input Guidelines:**
```python
# Good: Explicit dtype, contiguous array
arr = np.array([1.0, 2.0, 3.0], dtype=np.float64)
arr = np.ascontiguousarray(data)

# Good: Scalar inputs with consistent types
result = numba_func(1.0, 2.0)  # Both floats
result = numba_func(np.float64(1.0), np.float64(2.0))  # Explicit numpy scalars

# Avoid: Python objects, strings, dicts, sets, or lists as inputs
# Avoid: Object dtype arrays (dtype=object)
# Avoid: Structured arrays with complex nested types
```

**Common Pitfalls to Avoid in Tests:**
- Do not pass Python lists directly - convert to numpy arrays first
- Do not use `None` as default arguments in test inputs (unless the function explicitly handles it in nopython mode)
- Do not pass non-contiguous array views without checking if the function supports them
- Do not mix `int32` and `int64` across test cases for the same parameter

**Numba-Specific Assertions:**
```python
# Use numpy comparison functions
assert np.allclose(result, expected, rtol=1e-7, atol=1e-14)
assert np.array_equal(int_result, expected_int)  # For integer results
```

---

### PyTorch (`@torch.compile`)

**Supported Types:**
- `torch.Tensor` with any dtype (`float32`, `float64`, `int32`, `int64`, `bfloat16`, etc.)
- Python primitives (`int`, `float`, `bool`) - but prefer tensors
- Tuples and lists of tensors
- Dictionaries with string keys and tensor values

**Recompilation rules:** See the top-level "Recompilation Prevention" section above. All 6 rules apply to `torch.compile` — especially: one canonical shape, `float32` default, no rank mixing, no `None` inputs.

**Test Input Guidelines:**
```python
# GOOD: All tests use the same shape (4,) and dtype float32
def test_basic():
    x = torch.tensor([1.0, -2.0, 3.0, -4.0], dtype=torch.float32)
    result = my_func(x)
    assert torch.allclose(result, expected)

def test_zeros():
    x = torch.zeros(4, dtype=torch.float32)
    result = my_func(x)
    assert torch.allclose(result, expected)

def test_gradient():
    x = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, requires_grad=True)
    result = my_func(x)
    result.sum().backward()
    assert x.grad is not None

# GOOD: Edge cases isolated in one test
def test_edge_case_shapes():
    # These will each trigger a recompilation — keep them together
    assert my_func(torch.tensor([0.0])).shape == (1,)
    assert my_func(torch.randn(1000)).shape == (1000,)

# BAD: Every test uses a different shape — causes N recompilations
def test_small():
    x = torch.randn(3)       # shape (3,) — 1st compilation
def test_medium():
    x = torch.randn(100)     # shape (100,) — 2nd compilation (recompile!)
def test_large():
    x = torch.randn(10000)   # shape (10000,) — 3rd compilation (recompile!)
def test_2d():
    x = torch.randn(3, 4)    # rank 2 — 4th compilation (recompile!)
```

```python
# Good: Explicit dtype and device
x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
x = torch.randn(4, 4, dtype=torch.float32)

# Good: Consistent device placement
device = torch.device('cpu')  # or 'cuda' if available
x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device, dtype=torch.float32)

# Good: Using torch generator for reproducibility
g = torch.Generator().manual_seed(42)
x = torch.randn(4, 4, generator=g, dtype=torch.float32)

# Avoid: Mixing CPU and CUDA tensors in same operation
# Avoid: In-place operations on leaf tensors that require grad (if testing gradients)
# Avoid: Dynamic control flow that changes based on tensor values (causes graph breaks)
```

**Common Pitfalls to Avoid in Tests:**
- Do not use `tensor.item()` inside compiled functions - extract values outside
- Do not modify tensor shapes dynamically within compiled regions
- Do not use Python `print()` or logging inside compiled functions
- Do not rely on specific compilation behavior - test functional correctness only
- Do not pass `None` as a tensor argument — `torch.compile` will trace it and fail

**PyTorch-Specific Assertions:**
```python
# Use torch comparison functions
assert torch.allclose(result, expected, rtol=1e-5, atol=1e-8)
assert torch.equal(int_result, expected_int)  # For exact integer comparison

# For gradient testing
assert torch.allclose(tensor.grad, expected_grad, rtol=1e-4)
```

**Graph Break Considerations:**
When testing `@torch.compile` functions, avoid inputs that cause graph breaks:
- Data-dependent control flow
- Calls to non-compilable Python functions
- Dynamic shapes (unless using `dynamic=True`)

---

### TensorFlow (`@tf.function`)

**Supported Types:**
- `tf.Tensor` with any dtype
- `tf.Variable` (be cautious with mutations)
- Python primitives (traced as constants - be aware of retracing)
- `tf.TensorSpec` for input signatures
- Nested structures of tensors (lists, tuples, dicts)

**Test Input Guidelines:**
```python
# Good: Explicit dtype specification
x = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32)
x = tf.random.normal([100, 100], dtype=tf.float64)

# Good: Using input_signature to prevent retracing
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
def my_func(x):
    ...

# Good: Setting seeds for reproducibility
tf.random.set_seed(42)
x = tf.random.normal([100, 100])

# Avoid: Python objects that can't be converted to tensors
# Avoid: Changing Python argument values between calls (causes retracing)
# Avoid: Using tensor.numpy() inside @tf.function
```

**Common Pitfalls to Avoid in Tests:**
- Do not pass different Python primitive values across test cases (causes retracing)
- Do not use `tensor.numpy()` inside tf.function - only use outside
- Do not test with `tf.Variable` mutations unless the function is designed for it
- Do not rely on Python side effects inside tf.function (they only execute during tracing)

**TensorFlow-Specific Assertions:**
```python
# Use numpy conversion for assertions (outside tf.function)
result_np = result.numpy()
expected_np = expected.numpy()
np.testing.assert_allclose(result_np, expected_np, rtol=1e-5, atol=1e-8)

# Or use tf functions
tf.debugging.assert_near(result, expected, rtol=1e-5, atol=1e-8)

# For exact comparison
tf.debugging.assert_equal(int_result, expected_int)
```

**AutoGraph Considerations:**
- Python `if` statements are converted to `tf.cond` - ensure both branches have same output structure
- Python `for`/`while` loops are converted to `tf.while_loop` - ensure loop variables have consistent types
- Avoid `break`/`continue` in complex nested loops

---

### JAX (`@jax.jit`)

**Supported Types:**
- JAX arrays (`jax.Array`, `jnp.ndarray`)
- NumPy arrays (converted automatically)
- Python primitives (traced as tracers - use `static_argnums` for compile-time constants)
- PyTrees: nested structures of arrays (tuples, lists, dicts, namedtuples)

**Test Input Guidelines:**
```python
# Good: Using jax.numpy for array creation
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
x = jax.random.normal(jax.random.PRNGKey(42), (100, 100))

# Good: Explicit PRNG key handling (JAX requires explicit randomness)
key = jax.random.PRNGKey(42)
key, subkey = jax.random.split(key)
x = jax.random.normal(subkey, (100,))

# Good: Using static_argnums for non-traced arguments
@jax.jit
def func(x, axis):  # axis should be static if it's used in shape operations
    return jnp.sum(x, axis=axis)
# Mark 'axis' as static: jax.jit(func, static_argnums=(1,))

# Avoid: In-place mutations (JAX arrays are immutable)
# Avoid: Data-dependent shapes
# Avoid: Side effects (prints, global variable modifications)
```

**Common Pitfalls to Avoid in Tests:**
- Do not test with in-place array modifications - JAX arrays are immutable
- Do not pass different static argument values without recompiling
- Do not use Python random module - use `jax.random` with explicit keys
- Do not rely on side effects - JAX functions must be pure

**JAX-Specific Assertions:**
```python
# Use jax.numpy or numpy for comparisons
assert jnp.allclose(result, expected, rtol=1e-5, atol=1e-8)
assert jnp.array_equal(int_result, expected_int)

# For testing gradients
grad_fn = jax.grad(func)
computed_grad = grad_fn(x)
assert jnp.allclose(computed_grad, expected_grad, rtol=1e-4)
```

**Functional Purity Requirements:**
JAX requires functions to be pure (no side effects). In tests:
- Do not expect `print()` statements to execute during JIT-compiled calls
- Do not test functions that modify global state
- Ensure random operations use explicit `PRNGKey` arguments

**PyTree Handling:**
```python
# Good: Testing with PyTree inputs
inputs = {'a': jnp.array([1.0, 2.0]), 'b': jnp.array([3.0, 4.0])}
result = jitted_func(inputs)

# Good: Testing with nested structures
inputs = (jnp.array([1.0]), {'x': jnp.array([2.0, 3.0])})
result = jitted_func(inputs)
```

---

## Cross-Framework Compatibility Notes

When a function might be optimized with multiple JIT frameworks:

1. **Prefer numpy arrays as the common input format** - all frameworks can consume them
2. **Use framework-agnostic assertions when possible:**
   ```python
   # Convert to numpy for comparison
   result_np = np.asarray(result)  # Works for torch, tf, jax, numpy
   np.testing.assert_allclose(result_np, expected, rtol=1e-5)
   ```
3. **Use float32 for PyTorch, float64 for NumPy/Numba** — PyTorch defaults to float32 and many ops don't support integer dtypes under `torch.compile`. NumPy/Numba work best with float64
4. **Avoid framework-specific features in test inputs** unless testing that specific functionality
5. **Reuse the same tensor shape across all tests** to avoid recompilation overhead in JIT frameworks. Only vary shape/dtype in dedicated edge-case tests

In the case of Pytorch, TensorFlow, Jax and MLX -> Write CPU-only, CUDA-only and MPS-only test functions and skip if the device is not available.

---

## Summary: Inputs and Tests to AVOID for JIT Compatibility

To ensure tests work identically for both compiled and uncompiled code, **never** include:

### Inputs to Never Use
- `None` as function arguments
- Python lists (use numpy arrays instead)
- String arrays or object dtype arrays
- Mixed-type inputs across test cases
- Multidimensional arrays when function expects 1D (unless testing that specifically)

### Tests to Never Write
- Tests that use `pytest.raises()` for type/value errors
- Tests checking specific exception messages
- Tests for invalid input handling
- Tests that rely on Python-specific type coercion behavior

### Always Focus On
- Valid, typed numpy arrays or tensors with explicit dtypes
- Numerical correctness with tolerance-based comparisons (`allclose`)
- Edge cases that produce valid outputs (empty arrays returning empty results, single elements, etc.)
- Deterministic inputs using seeded RNGs
- Shape and dtype preservation
