[zephyr] Drop dead output_shard_count handling in non-scatter regroup

stage.output_shards is only ever set on Scatter (via group_by) and
Reshard stages, and Reshard short-circuits before reaching
_regroup_result_refs. The non-scatter branch's max-with-output_shard_count
was dead code suggesting semantics that don't exist; resharding belongs
to ReshardOp.

Author: ravwojdyla

lib/zephyr/src/zephyr/execution.py

@@ -1323,17 +1323,14 @@ def _regroup_result_refs(
     Each reducer reads the per-mapper ``.scatter_meta`` sidecars in parallel
     to build its own ``ScatterReader`` without coordinator-side consolidation.
     """
-    if is_scatter and output_shard_count is not None:
+    if is_scatter:
         # Scatter routes records into exactly ``output_shard_count`` buckets via
         # ``hash(key) % output_shard_count``; spawning more reduce tasks than that
         # produces empty output files for shard indices that no record hashes to.
-        num_output = output_shard_count
-    else:
-        num_output = max(max(result_refs.keys(), default=0) + 1, input_shard_count)
-        if output_shard_count is not None:
-            num_output = max(num_output, output_shard_count)
+        # When output_shard_count is None (group_by auto-detect), inherit the
+        # input shard count.
+        num_output = output_shard_count if output_shard_count is not None else input_shard_count
 
-    if is_scatter:
         # Collect all scatter file paths from all workers. The coordinator
         # does NOT read the sidecars or write a consolidated manifest —
         # reducers do their own parallel sidecar reads.
@@ -1344,7 +1341,9 @@ def _regroup_result_refs(
         shared_refs = MemChunk(items=all_paths)
         return [ListShard(refs=[shared_refs]) for _ in range(num_output)]
 
-    # Non-scatter: each result's shard maps to its own index
+    # Non-scatter: 1:1 mapping from input shard index to output. Resharding
+    # to a different shard count belongs to ReshardOp, not here.
+    num_output = max(max(result_refs.keys(), default=0) + 1, input_shard_count)
     return [result_refs[idx].shard if idx in result_refs else ListShard(refs=[]) for idx in range(num_output)]