refactor: use jnp.take_along_axis for return_last_only indexing

[stanley1208]

Summary

Resolves TODO(epot): Use jnp.take_along_axis in both Transformer.__call__ and Gemma3nTransformer.__call__.

The return_last_only code path previously used manual fancy indexing with jnp.arange to select the last non-padded token per batch element. This replaces it with jnp.take_along_axis, which is cleaner and avoids constructing an index array.

Changes

• gemma/gm/nn/_transformer.py (line ~253): replaced x[jnp.arange(len(x)), idx, ...] with jnp.take_along_axis(x, idx[:, None, None], axis=1) + squeeze
• gemma/gm/nn/gemma3n/_transformer.py (line ~323): same change

Test plan

• test_prefill passes (exercises return_last_only=True via the standard Transformer)
• test_sampler passes (end-to-end sampling using prefill)
• test_last_only passes (directly tests return_last_only on Gemma3n)
• All 9 relevant tests pass