Between master on Jul 18 and master on Oct 15, 2026, 15 of our PRs landed in Keras. Same GPU, same torch and jax versions, same models; only the Keras source changes. On the issue's test LM, compiled Keras[torch] now generates text 11.6× faster, torch.compile went from slowing Keras down to making it 6.56× faster than eager Keras, and an LLM forward makes 12–23% fewer Python calls on all four backends.
Generate-32 throughput, July master vs today
Tokens per second from the 2.1M-parameter test LM, 128 tokens per call (4 sequences × 32 steps); longer bars are faster. In July, compiling a Keras[torch] model made it slower: 1,458 tok/s compiled against 2,139 eager. Today the compiled model runs at 16,975 tok/s: 6.56× faster than eager Keras, within 1.53× of plain compiled PyTorch, down from 17.8×.
Three merges that made torch.compile work with Keras
Three merges each turned one compiled workload from slow to fast in a single step.
The three PRs: #23327 (attention), #23201 (dict ordering in the tree code) and #23647 (identity checks in the tree code). Graph breaks per forward pass, at every merge:
Less Python in every Keras call, on all four backends
Most of the series changed shared code (Layer.call, op tracebacks, dtype handling, input checks), so the saving reaches every backend, not just torch. Python function calls per forward pass, July to today: PyTorch −23% on the LM and −33% on the CNN; JAX −12% on the LM and −17% on the CNN; TensorFlow −13% on the LM and −19% on the CNN; NumPy −16% on the LM and −18% on the CNN. Adding #23336 takes NumPy to −36% on the LM. These counts are exact and repeat call for call on any machine.
Keras[jax]
Under jax.jit the Python runs once at trace time, so Keras[jax] was already close to plain JAX and stays there: 2.30 → 2.45 ms per generate-32 against 2.22 ms for plain JAX. Eager Keras[jax] is where the saving shows in the exact counts: −12% Python calls per forward. The timing agrees: the sweep gives 293 → 267 ms (1.10×) and the back-to-back A/B rounds give 1.08×.
Every config, every workload
Show all 12 configs and workloads (ms per call)| Keras config | Workload | Plain framework (ms) | Jul 18 (ms) | Oct 15 (ms) | Change | + #23336 (ms) | A/B check |
|---|---|---|---|---|---|---|---|
| Keras[torch] · eager | CNN forward | 0.064 | 0.668 | 0.477 | 1.40× | 0.466 | 1.39× |
| Keras[torch] · eager | LLM forward | 0.189 | 1.88 | 1.55 | 1.22× | 1.48 | 1.20× |
| Keras[torch] · eager | LLM generate-32 | 5.85 | 59.8 | 49.5 | 1.21× | 47.2 | 1.20× |
| Keras[torch] · torch.compile | CNN forward | 0.076 | 1.28 | 0.160 | 8.03× | 0.161 | 8.41× |
| Keras[torch] · torch.compile | LLM forward | 0.160 | 2.69 | 0.302 | 8.93× | 0.299 | 9.60× |
| Keras[torch] · torch.compile | LLM generate-32 | 4.92 | 87.8 | 7.54 | 11.6× | 7.66 | 12.4× |
| Keras[jax] · eager | CNN forward | 0.508 | 1.34 | 1.23 | 1.08× | 1.18 | 1.15× |
| Keras[jax] · eager | LLM forward | 2.79 | 8.46 | 8.38 | 1.01× | 7.90 | 1.12× |
| Keras[jax] · eager | LLM generate-32 | 89.0 | 293 | 267 | 1.10× | 269 | 1.08× |
| Keras[jax] · jax.jit | CNN forward | 0.077 | 0.079 | 0.081 | 0.97× | 0.080 | 1.02× |
| Keras[jax] · jax.jit | LLM forward | 0.120 | 0.131 | 0.135 | 0.97× | 0.130 | 1.05× |
| Keras[jax] · jax.jit | LLM generate-32 | 2.22 | 2.30 | 2.45 | 0.94× | 2.27 | 1.05× |
A/B check: July ÷ today from 3 separate back-to-back rounds (higher is a bigger speedup).
Each of our PRs, measured where it landed
Each PR was timed on the commit just before and just after it merged, so every row is that one PR's own effect. Δ is just after the merge minus just before it; negative is better.
Show the per-PR table| PR | Merged | What it changed | Δ calls/fwd torch | JAX | TF | NumPy | Δ eager gen-32 (ms) | Δ compiled gen-32 (ms) | Graph breaks |
|---|---|---|---|---|---|---|---|---|---|
| #23183 | 07-28 | Traceback details are built only when an op fails, not on every call | −598 | −586 | −622 | −546 | −2.9 | −5.2 | · |
| #23327 | 08-30 | Removed the torch.compile graph break in the torch backend's dot_product_attention | +6 | 0 | 0 | 0 | +0.40 | −12.3 | LLM 11 → 0 |
| #23189 | 09-03 | Fixed an always-true device check that sent every tensor through .to() | 0 | +38 | 0 | 0 | +0.28 | −1.1 | · |
| #23200 | 09-05 | Layer.call stops re-setting built/_called on every call | −1,848 | −1,098 | −1,188 | −1,100 | −4.9 | −10.5 | · |
| #23182 | 09-10 | O(1) dtype membership and a fast path for canonical dtype strings | −300 | −82 | −497 | −116 | −0.89 | −1.5 | · |
| #23199 | 09-10 | Dropped a redundant channels_last copy on torch conv inputs | 0 | −50 | 0 | 0 | +0.40 | −0.29 | · |
| #23201 | 09-11 | Torch tree code made traceable: canonical dict key order | −142 | +8 | 0 | 0 | +0.55 | −5.2 | CNN 9 → 0 |
| #23296 | 09-11 | Keeps call-context kwargs on the recorded Node (fix) | −109 | −78 | −14 | −46 | −0.24 | −0.09 | · |
| #23187 | 09-12 | TorchTrainer skips a full flatten walk to find the batch size | 0 | −26 | 0 | 0 | +0.14 | −0.21 | · |
| #23301 | 09-12 | Layer name scope opened once per call, not twice | −330 | −568 | −638 | −330 | −0.89 | −0.21 | · |
| #23190 | 09-15 | torch slice() keeps dynamic dims under torch.export | 0 | −16 | 0 | 0 | −0.23 | −0.16 | · |
| #23647 | 09-30 | Torch tree code stops comparing id()s | 0 | −8 | 0 | 0 | −0.58 | −49.4 | · |
| #23645 | 10-07 | torch pooling ops run in channels_last | 0 | −16 | 0 | 0 | +0.10 | −0.05 | · |
| #23184 | 10-13 | Fast path for single-InputSpec input checks | −102 | −59 | −63 | −63 | +0.06 | +0.05 | · |
| #23920 | 10-15 | Functional models reject a mismatched last axis (fix) | 0 | −54 | 0 | 0 | +0.10 | +0.02 | · |
| #23336 | not merged yet (approved) | Memoized dtype promotion in result_type | −713 | −378 | −1,471 | −3,399 | −2.2 | +0.12 | · |
Hot functions in one Keras[torch] LLM forward
Show the hot-functions table| Function | Jul 18 | Oct 15 | Change | + #23336 |
|---|---|---|---|---|
| isinstance | 2,600 | 1,740 | −33% | 1,680 |
| traceback error wrappers | 74 | 48 | −35% | 48 |
| thread-local config lookups | 624 | 513 | −18% | 513 |
| standardize_dtype | 177 | 244 | +38% | 184 |
| tree.flatten | 93 | 67 | −28% | 67 |
| tree.map_structure | 102 | 62 | −39% | 62 |
| InputSpec checks | 14 | 14 | +0% | 14 |
| CallSpec objects built | 23 | 23 | +0% | 23 |
| inspect.Signature.bind | 22 | 22 | +0% | 22 |
| convert_to_tensor | 86 | 86 | +0% | 86 |
| get_device | 138 | 138 | +0% | 138 |
How this was measured
- Machine: one NVIDIA RTX PRO 6000 Blackwell, kept idle during the runs.
- Software: the same for every run: torch 2.13.0, jax 0.11.2, Python 3.14.7. Only the Keras source changes.
- Commits: 31 snapshots of master: Jul 18 (
a2f8eac8), just before and just after each of the 15 merges, Oct 15 (cc9de2b31), and Oct 15 plus #23336. - Models: the issue's two test models: a small CNN (114k parameters) and a 2-layer transformer LM (2.1M parameters, 32-token greedy generate).
- Timing: the issue's benchmark script, one process at a time, warm-up first, a GPU sync after every call. Each number is a median, and the whole sweep ran three times for PyTorch and twice for JAX.
- Cross-check: a separate back-to-back run of July against October gives 12.4× compiled and 1.20× eager (the sweep: 11.6× and 1.21×).
- Counts: Python calls from cProfile, graph breaks from
torch._dynamo.explain. These are exact and repeat on any machine. - Not covered: training, CPU, larger models, and TensorFlow or NumPy timing (call counts only).
The work is tracked in issue #22561. Thank you to the Keras maintainers who reviewed and merged these PRs.





