# UNDF: UNDF-2026-000000431
# jax-0001 — pallas/fuser/jaxpr_fusion.py O(G × I × F) kernel fusion membership test

## Location

`jax/_src/pallas/fuser/jaxpr_fusion.py`, lines 211–216

## Defective Code

```python
# line 211 — outer loop over G fusion output groups
for i, outvars_group in enumerate(partial_flat):
    flat_group_vars, _ = tree_util.tree_flatten(outvars_group)   # list of size F

    # line 215 — inner loop over I jaxpr inputs; flat_group_vars is a list → O(F) membership
    used_jaxpr_invars = [False] * len(all_values) + [
        v in flat_group_vars for v in jaxpr_out.invars   # O(I × F) per group
    ]
```

`flat_group_vars` is a Python `list`. The comprehension `v in flat_group_vars` performs a
linear scan of F elements for each of I `jaxpr_out.invars` entries.
With G fusion groups this is **O(G × I × F)** total.

The same pattern also appears at lines 186–188 in the preceding loop, though that uses a
`set` (`used_fusible_outvars`), so it is already O(1) — only lines 215–216 are defective.

## Fixed Code

```python
for i, outvars_group in enumerate(partial_flat):
    flat_group_vars, _ = tree_util.tree_flatten(outvars_group)
    flat_group_vars_set = set(flat_group_vars)          # O(F) — build once per group

    used_jaxpr_invars = [False] * len(all_values) + [
        v in flat_group_vars_set for v in jaxpr_out.invars   # O(1) each → O(I) total
    ]
```

## Complexity Analysis

| Path | Complexity |
|------|-----------|
| Slow (original) | O(G × I × F) |
| Fast (fixed) | O(G × (I + F)) |

For G=8 output groups, I=64 jaxpr inputs, F=32 group vars: **32× speedup**.
For large Pallas GPU kernels with many fused outputs: up to **100× speedup**.

## Severity

**MEDIUM** — executed once per Pallas kernel compilation via `jaxpr_fusion.fuse_jaxpr`.
Not on the per-step inference hot path, but drives compile time for large Pallas GPU kernels
(attention layers, Flash Attention variants). Slow compilation blocks model iteration.
