Refactor SamplerLoop class into standalone functions

Current Behavior
gemma/gm/text/_sampler_loop.py
currently implements sampling as a
SamplerLoop
class (line 102-266). This class only contains static, hashable attributes and serves primarily as a namespace for related functions, which is not idiomatic Python for this use case.

There is a TODO at line 98-100:

python
# TODO(epot): Refactor into simple function, rather than class.
# * autoregressive_sample()
# * autoregressive_stream_sample()
The current class-based approach:

Adds unnecessary boilerplate (dataclass definition, self parameters)
Makes the code harder to understand (class usually implies stateful behavior)
Reduces functional programming benefits in a JAX-heavy codebase
Expected Behavior
The sampling logic should be refactored into two standalone, pure functions that better align with JAX's functional programming paradigm:

autoregressive_sample() - Non-streaming sampling
autoregressive_stream_sample() - Streaming sampling
This would simplify the API, improve code clarity, and make it easier to compose with other JAX transformations.

Proposed Solution
I will refactor _sampler_loop.py by:
Extract configuration parameters into a dedicated dataclass (e.g., SamplerConfig)
Convert SamplerLoop.sample() into autoregressive_sample() standalone function
Convert SamplerLoop._sample_loop() into internal helper or merge with main function
Convert SamplerLoop._stream_sample_loop() into autoregressive_stream_sample() standalone function
Update all call sites in
_sampler.py
_chat_sampler.py
, and
_tool_sampler.py
to use the new function-based API
Update tests in
_sampler_test.py
and
sampler_e2e_test.py
Environment
Gemma version: main branch
Python: 3.x (JAX-compatible)