jax-ml/jax
Languages
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
Harness | Input / Output Cost | ||||||
|---|---|---|---|---|---|---|---|
1 |  Mini-SWE-agent | 43 / 54 | $2.27 | $5/$30 | 10m11s | ||
2 | Mini-SWE-agent | 42 / 54 | $1.05 | $2/$6 | 6m35s | ||
3 | Mini-SWE-agent | 42 / 54 | $2.55 | $5/$30 | 7m49s | ||
4 | Mini-SWE-agent | 42 / 54 | $10.42 | $3/$15 | 18m24s | ||
5 | Mini-SWE-agent | 42 / 54 | $7.26 | $10/$50 | 18m41s | ||
6 | Mini-SWE-agent | 42 / 54 | $4.31 | $3/$15 | 24m52s | ||
7 | Mini-SWE-agent | 41 / 54 | $0.49 | $2/$12 | 3m07s | ||
8 | Mini-SWE-agent | 40 / 54 | $0.15 | $0.2/$1.2 | 6m22s | ||
9 | Mini-SWE-agent | 40 / 54 | $1.56 | $1.5/$9 | 6m34s | ||
10 | Mini-SWE-agent | 40 / 54 | $4.61 | $5/$25 | 15m28s | ||
11 | Mini-SWE-agent | 40 / 54 | $3.89 | $5/$25 | 16m08s | ||
12 | Mini-SWE-agent | 38 / 54 | $1.04 | $1.4/$4.4 | 18m50s | ||
13 | Mini-SWE-agent | 37 / 54 | $1.10 | $2/$12 | 5m21s | ||
14 | Mini-SWE-agent | 33 / 54 | $0.45 | $1/$5 | 5m02s |
Key Takeaways
- Claude Fable 5, Claude Sonnet 5, Kimi K3, Grok 4.5, and GPT 5.5 with Mini-SWE-agent each resolve 42 of 54 tasks, one fewer than GPT-5.6 Sol.
- GPT-5.6 Terra with Mini-SWE-agent resolves 41 of 54 tasks, while GPT-5.6 Luna and Gemini 3.5 Flash each resolve 40 of 54.
- All supplied results cover 54 tasks, so these are not smoke tests or directional results.
Model Comparison
Accuracy
79.63%
GPT-5.6 Sol
77.78%
Grok 4.5
Task outcomes
54 tasks
Cost / test
$2.27
GPT-5.6 Sol
$1.05
Grok 4.5
Cost distribution
Latency
10m 11s
GPT-5.6 Sol
6m 35s
Grok 4.5
Latency distribution
Cost Analysis
Average Token Use / Test
Cost is the clearest tradeoff in this comparison. GPT-5.6 Sol leads at 79.63% for $2.27 per test. GPT-5.6 Terra is the lower-cost option at 75.93% for $0.49 per test.
Latency Analysis
Average Response Time / Test
Latency separates several models with similarly strong scores. GPT-5.6 Sol leads at 79.63%, while GPT-5.6 Terra is fastest at 3m 7s with 75.93% accuracy.
Tasks with failures
| Models | |||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| GPT-5.6 Sol | |||||||||||||||||||||||||||
| Grok 4.5 | |||||||||||||||||||||||||||
| GPT 5.5 | |||||||||||||||||||||||||||
| Claude Sonnet 5 | |||||||||||||||||||||||||||
| Claude Fable 5 | |||||||||||||||||||||||||||
| Kimi K3 | |||||||||||||||||||||||||||
| GPT-5.6 Terra | |||||||||||||||||||||||||||
| GPT-5.6 Luna | |||||||||||||||||||||||||||
| Gemini 3.5 Flash | |||||||||||||||||||||||||||
| Claude Opus 4.8 | |||||||||||||||||||||||||||
| Claude Opus 4.7 | |||||||||||||||||||||||||||
| GLM 5.2 | |||||||||||||||||||||||||||
| Gemini 3.1 Pro Preview (02/26) | |||||||||||||||||||||||||||
| Claude Haiku 4.5 (Nonthinking) |
Task detail
deabe65Issue statement
When using shard_map (jax.shard_map) with out_specs that declare one or more mesh axes as unreduced (for example out_specs=P('x', None, unreduced={'y'})), shard_map currently inspects the values returned by the mapped function and, when a returned value is varying over an axis that out_specs says should be unreduced, it silently inserts a varying->unreduced cast so that the output matches out_specs.
This implicit varying->unreduced conversion is surprising and undesirable: producing an unreduced output is a semantically meaningful collective-like step, and hiding it behind out_specs makes the program's behavior non-obvious. shard_map should not perform this implicit cast on the user's behalf.
Desired behavior:
shard_mapshould still implicitly broadcast/pvarya value to become varying over axes required byout_specs(this partial-vary behavior is unchanged).shard_mapshould NOT implicitly cast a varying value into anunreducedvalue to satisfyout_specs. If the mapped function returns a value that is not alreadyunreducedalong an axis thatout_specsmarks asunreduced,shard_mapshould raise the standardout_specsmismatch error (aValueErrorwhose message refers to theout_specspassed toshard_map) rather than silently inserting the cast.- To obtain an
unreducedoutput, the user must explicitly cast inside the mapped function, e.g.c = jax.lax.pcast(c, 'y', to='unreduced'), which makes the valueunreducedover'y'(and no longer varying over'y'). After such an explicit cast, the value satisfies anunreduced={'y'}out_specsand the call succeeds as before.
View Hidden Tests
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.pyindex 68f2a6980..fd3bcb445 100644--- a/tests/shard_map_test.py+++ b/tests/shard_map_test.py@@ -222,6 +222,9 @@ class ShardMapTest(jtu.JaxTestCase): def f(a, b): c = jnp.einsum('ab,bc->ac', a, b) self.assertEqual(c.aval.mat.varying, {'x', 'y'})+ c = jax.lax.pcast(c, 'y', to='unreduced')+ self.assertEqual(c.aval.mat.varying, {'x'})+ self.assertEqual(c.aval.mat.unreduced, {'y'}) return c out = f(arr1, arr2)@@ -262,6 +265,8 @@ class ShardMapTest(jtu.JaxTestCase): def f(a, b): c = jnp.dot(a, b) self.assertEqual(c.aval.mat.varying, {'x'})+ c = jax.lax.pcast(c, 'x', to='unreduced')+ self.assertEqual(c.aval.mat.unreduced, {'x'}) return c out = f(arr1, arr2)@@ -282,9 +287,7 @@ class ShardMapTest(jtu.JaxTestCase): self.assertEqual(c.aval.mat.varying, {'x'}) return c - with self.assertRaisesRegex(- ValueError,- "vary_unreduced_cast is a Varying->Unreduced collective"):+ with self.assertRaisesRegex(ValueError, "out_specs passed to shard_map"): f(arr1, arr2) def test_matmul_reduce_scatter(self):@@ -5075,7 +5078,7 @@ class ShardMapTest(jtu.JaxTestCase): @jax.jit @shard_map(in_specs=P('x'), out_specs=P(unreduced={'x'})) def f(x):- return x+ return jax.lax.pcast(x, 'x', to='unreduced') out = f(arr) self.assertTupleEqual(out.shape, (2,))