Restore previous _compute_stats behavior in local case.
Some earlier code refactors introduced a concat+split "no-op" to the non-distributed codepath
that turns out not to be such a no-op and can cause issues in some compiled code. We polish
the refactor a bit to avoid this introduced change.
PiperOrigin-RevId: 555039226
