Fix incorrect batch-dim shape annotations on vmapped _get_new_text_tokens_inner

[PratikDhanave]

What

In gemma/gm/vision/_token_utils.py, _get_new_text_tokens_inner is called inside a vmap:

return jax.vmap(_get_new_text_tokens_inner, in_axes=(0, 0, None, None))(
    mm_start, text_tokens, offset_by, length_with_mm
)

in_axes=(0, 0, None, None) strips the leading axis, so the callee operates on single rows of shape [L]. But its parameters were annotated as batched:

def _get_new_text_tokens_inner(
    mm_start: Bool['B L'],   # <- wrong, should be ['L']
    text_tokens: Int['B L'], # <- wrong, should be ['L']
    ...
) -> Int['L']:

This contradicts three things:

• the function's own docstring — "_get_new_text_tokens_positions without batch dimension" — and its -> Int['L'] return type;
• the vmap in_axes that remove axis 0;
• the sibling _get_new_mm_tokens_inner, which correctly annotates its post-vmap arg as Bool['L'].

The body confirms rank-1: it allocates jnp.zeros((length_with_mm,)) and indexes it 1-D.

Change

Correct the two annotations to ['L'], matching the sibling's # pyrefly: ignore[not-a-type, unknown-name] style. The function is not @typechecked, so this is an annotation/static-analysis fix with no runtime behavior change.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.