Fix pytype late-directive warnings.

PiperOrigin-RevId: 949750872

Author: matthiaskramm

src/maxtext/inference/maxengine/maxengine.py

@@ -1099,8 +1099,8 @@ def _prefill_multisampling_jit(
       )
       first_generated_tokens.append(first_generated_token)
       if self.config.return_log_prob:
-        # pytype: disable=attribute-error
-        token_logps.append(inference_utils.log_prob_of_chosen_token(selected_logits, first_generated_token))
+        p = inference_utils.log_prob_of_chosen_token(selected_logits, first_generated_token)
+        token_logps.append(p)  # pytype: disable=attribute-error
     first_generated_tokens = jnp.concatenate(first_generated_tokens, axis=0)
     if self.config.return_log_prob:
       token_logps = jnp.concatenate(token_logps, axis=0)

src/maxtext/models/qwen3.py

@@ -658,14 +658,12 @@ def __call__(
       # vLLM PAGED STATE PATH: use tpu_inference fused conv + ragged delta-rule.
       # =========================================================================
       try:
-        # pylint: disable=import-outside-toplevel
-        # pytype: disable=import-error
-        from tpu_inference.layers.common.gdn_attention import GdnAttentionConfig, run_jax_gdn_attention  # pylint: disable=import-outside-toplevel
-        from tpu_inference.layers.common.ragged_gated_delta_rule_wrapper import RaggedGatedDeltaRuleImpl  # pylint: disable=import-outside-toplevel
-        from tpu_inference.layers.common.sharding import ShardingAxisName  # pylint: disable=import-outside-toplevel
-        from tpu_inference.layers.common.utils import reorder_concatenated_tensor_for_sharding  # pylint: disable=import-outside-toplevel
-        from tpu_inference.utils import get_mesh_shape_product  # pylint: disable=import-outside-toplevel
-        from jax.sharding import PartitionSpec as P_spec  # pylint: disable=import-outside-toplevel
+        from tpu_inference.layers.common.gdn_attention import GdnAttentionConfig, run_jax_gdn_attention  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+        from tpu_inference.layers.common.ragged_gated_delta_rule_wrapper import RaggedGatedDeltaRuleImpl  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+        from tpu_inference.layers.common.sharding import ShardingAxisName  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+        from tpu_inference.layers.common.utils import reorder_concatenated_tensor_for_sharding  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+        from tpu_inference.utils import get_mesh_shape_product  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
+        from jax.sharding import PartitionSpec as P_spec  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
       except ImportError as e:
         raise ImportError(
             "GDN attention kernel require the vllm-tpu package. Please install it with `pip install vllm-tpu`."