[flax:examples:wmt] Small linter fixes.

PiperOrigin-RevId: 838880356

Author: danielsuo

examples/wmt/models.py

@@ -20,8 +20,8 @@
 # pytype: disable=wrong-keyword-args
 # pytype: disable=attribute-error
 
-from typing import Any, Optional
 from collections.abc import Callable
+from typing import Any
 
 from flax import linen as nn
 from flax import struct
@@ -549,7 +549,8 @@ def decode(
 
     # Make padding attention masks.
     if config.decode:
-      # for fast autoregressive decoding only a special encoder-decoder mask is used
+      # for fast autoregressive decoding only a special encoder-decoder mask is
+      # used
       decoder_mask = None
       encoder_decoder_mask = nn.make_attention_mask(
           jnp.ones_like(targets) > 0, inputs > 0, dtype=config.dtype

examples/wmt/train.py

@@ -250,7 +250,7 @@ def loss_fn(params):
   if state.dynamic_scale:
     # if is_fin == False the gradients contain Inf/NaNs and optimizer state and
     # params should be restored (= skip this step).
-    select_fn = functools.partial(jnp.where, is_fin)
+    select_fn = functools.partial(jnp.where, is_fin)  # pylint: disable=undefined-variable
     new_state = new_state.replace(
         opt_state=jax.tree_util.tree_map(
             select_fn, new_state.opt_state, state.opt_state
@@ -259,7 +259,7 @@ def loss_fn(params):
             select_fn, new_state.params, state.params
         ),
     )
-    metrics["loss_scale"] = dynamic_scale.scale * metrics["denominator"]
+    metrics["loss_scale"] = dynamic_scale.scale * metrics["denominator"]  # pylint: disable=undefined-variable
 
   return new_state, metrics
 
@@ -649,8 +649,8 @@ def decode_tokens(toks):
           metrics_sums = jax.tree_util.tree_map(jnp.sum, train_metrics)
           denominator = metrics_sums.pop("denominator")
           summary = jax.tree_util.tree_map(
-              lambda x: x / denominator, metrics_sums
-          )  # pylint: disable=cell-var-from-loop
+              lambda x: x / denominator, metrics_sums  # pylint: disable=cell-var-from-loop
+          )
           summary["learning_rate"] = lr
           summary = {"train_" + k: v for k, v in summary.items()}
           writer.write_scalars(step, summary)