Type mismatch: _tokenize_prompts() doesn't accept tuple pad_length but Sampler.pad_length defaults to tuple

# Type mismatch: `_tokenize_prompts()` doesn't accept tuple `pad_length` but `Sampler.pad_length` defaults to tuple

## Problem

There's a type signature mismatch between `Sampler.pad_length` and `_tokenize_prompts()` that causes a `TypeError` when a tuple is passed. The `Sampler` class defaults `pad_length` to `(256, 512, 1024)` which is a tuple but `_tokenize_prompts()` only accepts `int | None` and doesn't handle tuples properly.

Looking at `gemma/gm/text/_sampler.py` shows `pad_length: None | int | tuple[int, ...] = (256, 512, 1024)` but the method signature only accepts `pad_length: int | None = None` and the logic uses `max_prompt_len = pad_length or max(len(t) for t in tokens)` which breaks when `pad_length` is a tuple because tuples are truthy so the `or` operator returns the tuple itself instead of calculating max length.

## Error Details

When calling `_tokenize_prompts()` directly with the tuple `pad_length` you get `TypeError: '>' not supported between instances of 'int' and 'tuple'`. The traceback shows the error originates from `_functional.py` where it tries to compare `seq_length > max_length` but `max_length` is a tuple `(256, 512, 1024)` instead of an integer. The call stack goes through `_tokenize_prompts()` calling `_functional.pad(tokens, max_length=max_prompt_len)` which then calls the inner `_pad()` function and that's where the comparison fails because `seq_length` is an int but `max_length` ended up being a tuple.

The issue is that when `pad_length` is `(256, 512, 1024)` the line `max_prompt_len = pad_length or max(len(t) for t in tokens)` evaluates to the tuple itself since tuples are truthy in Python so the `or` operator short-circuits and returns the first truthy value which is the tuple. Then this tuple gets passed down to `_functional.pad()` as `max_length` and when it tries to do `if seq_length > max_length:` Python can't compare an int to a tuple so it raises TypeError.

## Why This Matters

Even though the bug doesn't manifest in normal usage because `_get_inputs()` doesn't pass `pad_length` to `_tokenize_prompts()` this is still a real bug that needs fixing. The type signatures are inconsistent which creates API confusion and if someone refactors the code to pass `pad_length` or calls `_tokenize_prompts()` directly with `sampler.pad_length` it will crash immediately. The `_prefill.prefill()` method handles tuples correctly by using `_pad_to_bucket()` to pick smallest bucket that fits so there's an inconsistency in how different methods handle the same parameter.

## Solution

The fix updates the type signature and replaces the buggy logic:

1. Changed `pad_length: int | None = None` to `pad_length: int | tuple[int, ...] | None = None`
2. Replaced `max_prompt_len = pad_length or max(len(t) for t in tokens)` with proper handling for all three cases:
- `None`: use actual max length
- `int`: use that value directly
- `tuple`: iterate through bucket sizes and pick the smallest one that fits

This makes `_tokenize_prompts()` consistent with `_prefill.prefill()` which already handles tuples correctly.

Since this is a simple fix I'll submit a PR immediately after this issue is created.