[*.py] Resolve pylint c-class: C0301,C3001,C0114,C0115,C0116,C0200,C0121,C0201,C0206,C0209,C0412,C0415,C2801 ; [code_style.sh,.github/workflows/CPUTests.yml] Enable c-class

Author: SamuelMarks

.github/workflows/CPUTests.yml

@@ -36,7 +36,7 @@ jobs:
         pytype --jobs auto --disable 'import-error,late-directive,wrong-arg-types,module-attr,unsupported-operands' MaxText/ || true
     - name: Analysing the code with pylint in Maxtext/
       run: |
-         pylint --disable C0301,C3001,C0114,C0115,C0116,C0200,C0121,C0201,C0206,C0209,C0412,C0415,C2801,E0102,E0606,E1102,E1111,E1123,E1135,E1136,R0401,R1701,R1703,R1710,R1711,R1735,R0917,R1714,R1716,R1719,R1721,R1728,R1728,W0102,W0107,W0201,W0212,W0221,W0237,W0404,W0611,W0612,W0613,W0621,W0622,W0631,W0707,W0718,W1201,W1203,W1309,W1514,W4901 MaxText/ && \
+         pylint --verbose --msg-template='[{abspath}] {msg_id}:{line:3d},{column}: {obj}: {msg}' --disable E0102,E0606,E0611,E1102,E1111,E1120,E1121,E1123,E1135,E1136,R0401,R1701,R1703,R1710,R1711,R1735,R0917,R1714,R1716,R1719,R1721,R1728,R1728,W0102,W0107,W0201,W0212,W0221,W0223,W0237,W0404,W0611,W0612,W0613,W0621,W0622,W0631,W0707,W0718,W1201,W1203,W1309,W1514,W4901 MaxText/ && \
          echo 'Maxtext PyLint check successful' || { echo \
          'PyLint check has failed. Please run bash code_style.sh to fix issues'; exit 20; }
     - name: Analysing the code with pylint in pedagogical_examples/

MaxText/common_types.py

@@ -71,6 +71,8 @@
 
 
 class DecoderBlockType(enum.Enum):
+  """Decoder block types."""
+
   DEFAULT = "default"
   LLAMA2 = "llama2"
   MISTRAL = "mistral"

MaxText/convert_gpt3_ckpt_from_paxml.py

@@ -11,8 +11,6 @@
 limitations under the License.
 """
 
-from MaxText.globals import PKG_DIR
-
 # pylint: disable=line-too-long
 """Convert weights from a paxml gpt3 model to a MaxText one.
 
@@ -34,27 +32,31 @@
   --run-name=$RUN_NAME \
   --base-output-directory=$BASE_OUTPUT_DIR
 """
-from MaxText import max_utils
-from MaxText import maxtext_utils
-from MaxText import optimizers
-from MaxText import pyconfig
+
+import argparse
+import gc
 import os
+import sys
+
+from psutil import Process
+
+import numpy as np
+
+import jax
 from jax import random
 from jax.sharding import Mesh
-from MaxText.layers.models import Transformer
-from MaxText.layers import quantizations
-from MaxText import checkpointing
 
-import numpy as np
 import tensorstore as ts
 
-import sys
-import jax
-import gc
+from MaxText import checkpointing
 from MaxText import max_logging
-from psutil import Process
+from MaxText import maxtext_utils
+from MaxText import optimizers
+from MaxText import pyconfig
+from MaxText.globals import PKG_DIR
+from MaxText.layers import quantizations
+from MaxText.layers.models import Transformer
 from MaxText.train import save_checkpoint
-import argparse
 
 
 def fmt_size(num_bytes: int) -> str:

MaxText/decode.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CLI utility for running inference on a single/multi stream(s)"""
+"""CLI utility for running inference on a single/multi stream(s)."""
 
 import os
 from typing import Sequence
@@ -65,7 +65,7 @@ def _pad_to_batch_size(data: jax.Array, batch_size: int):
   )
 
   def _all_equals(elements: Sequence[jax.Array], target: jax.Array):
-    """Checks if each element equals the given target"""
+    """Checks if each element equals the given target."""
     stacked = jnp.stack(elements)
     row_comparisons = stacked == target
     return jnp.all(row_comparisons)

MaxText/deepseek_fp8_to_bf16.py

@@ -70,7 +70,7 @@ def weight_dequant_cpu(x: torch.Tensor, s: torch.Tensor, block_size: int = 128)
   return y
 
 
-def convert_fp8_to_bf16(fp8_path: string, bf16_path: string, cache_file_num: int = 2):
+def convert_fp8_to_bf16(fp8_path: str, bf16_path: str, cache_file_num: int = 2):
   """
   Converts a FP8 model to a BF16 model and saves the converted weights.
 

MaxText/elastic_train.py

@@ -187,7 +187,7 @@ def train_loop(config, elastic_manager, state=None):
   Args:
     config:
     state:
-    ckpt_path:
+    elastic_manager:
   Returns:
   """
   # Create a GoodputRecorder to log information

MaxText/inference/kvcache.py

@@ -16,13 +16,16 @@
 
 from typing import Any, Optional, Tuple
 
+import jax
+import jax.numpy as jnp
+
+from flax import linen as nn
+
 from aqt.jax.v2 import aqt_tensor
 from aqt.jax.v2 import config as aqt_config
 from aqt.jax.v2.flax import aqt_flax
+
 from MaxText import common_types
-from flax import linen as nn
-import jax
-import jax.numpy as jnp
 
 Array = common_types.Array
 AxisNames = common_types.AxisNames
@@ -104,6 +107,7 @@ def einsum_fn_with_rhs_qtensor(
       lhs_dequant_mode=None,
       lhs_calibration_mode=None,
   ):
+    """einsum function where QTensor is the right-hand-side"""
     # Assumes kv is already quantized.
     einsum = jnp.einsum
     if isinstance(kv, aqt_tensor.QTensor):
@@ -141,6 +145,7 @@ def einsum_fn_with_rhs_qtensor(
     return einsum
 
   def einsum_fn_with_rhs_qtensor_and_dequant(self, value):
+    """Get einstein summation for different dequant modes."""
     if self.dtype == jnp.float8_e4m3fn:
       return self.einsum_fn_with_rhs_qtensor(
           value,
@@ -184,6 +189,7 @@ def _get_cache_scale_logical_shape(self, batch, heads, cache_length):
     raise f"Invalid config for kv_quant_axis:{self.kv_quant.axis_cfg}"
 
   def _get_prefill_cache_vars(self, batch, key_heads, value_heads, key_head_size, value_head_size, model_mode):
+    """Get a shaped abstraction of the state"""
 
     cache_length = self.max_prefill_length
     dtype = self._get_cached_kv_dtype()
@@ -258,6 +264,7 @@ def _get_prefill_cache_vars(self, batch, key_heads, value_heads, key_head_size,
     return key_vars, value_vars, cached_segment_id_var
 
   def _get_ar_cache_vars(self, batch, key_heads, value_heads, key_head_size, value_head_size, model_mode):
+    """get ar cache vars"""
 
     dtype = self._get_cached_kv_dtype()
     if self.max_target_length <= self.max_prefill_length:
@@ -602,6 +609,7 @@ def value_body(i, val):
       )
 
   def get_cached_values(self, cache_vars, target_dtype, cache_axis_order) -> jax.Array | KVTensor:
+    """get cached values"""
     cache_var, cache_scale_var = cache_vars
     cache_value = cache_var.value
     if cache_scale_var is not None:

MaxText/inference/page_manager.py

@@ -23,10 +23,13 @@
 """
 
 from functools import partial
+from typing import Tuple
+
 import jax
 import jax.numpy as jnp
+
 from flax import struct
-from typing import Tuple
+
 from jaxtyping import Array, Integer, Bool
 
 from MaxText import common_types
@@ -477,7 +480,8 @@ def _validate_init_params(self) -> None:
     min_required = (self.max_target_length + self.tokens_per_page - 1) // self.tokens_per_page
     if self.max_pages_per_group < min_required:
       raise ValueError(
-          f"`pagedattn_max_pages_per_group` ({self.max_pages_per_group}) is insufficient for `max_target_length` ({self.max_target_length}). Needs {min_required}."
+          f"`pagedattn_max_pages_per_group` ({self.max_pages_per_group}) is insufficient for `max_target_length` "
+          f"({self.max_target_length}). Needs {min_required}."
       )
     # Check > 1 due to potential page 0 workaround
     if self.num_pages <= 1:

MaxText/inference/paged_attention.py

@@ -52,6 +52,8 @@
 
 
 class PagedAttentionOp(nn.Module):
+  """paged-attention op"""
+
   mesh: Mesh
   num_pages: int
   tokens_per_page: int
@@ -101,6 +103,7 @@ def init_or_get_kv_pages(self, model_mode: str):
     return key_pages_var, value_pages_var
 
   def paged_dot_product_attention_with_max_and_sum(self, query, key, value):
+    """paged dot product attention with max & sum"""
     b, t, n, d = query.shape
     _, s, n_kv, _ = key.shape
     query = jnp.reshape(query, (b, t, n_kv, n // n_kv, d))
@@ -307,6 +310,7 @@ def update_prefill_step_pages(
     value_pages_var.value = nn.with_logical_constraint(value, self.kv_pages_axis_names)
 
   def update_decode_step_pages(self, key_pages_var, value_pages_var, key, value, page_state):
+    """Update decode-step pages"""
     key_pages = key_pages_var.value
     value_pages = value_pages_var.value
 

MaxText/inference/paged_attention_kernel_v2.py

@@ -79,6 +79,7 @@ def ref_ragged_paged_attention(
     sm_scale: float = 1.0,
     mask_value: float = DEFAULT_MASK_VALUE,
 ):
+  """Ref ragged paged attention."""
   _, _, num_kv_heads, head_dim = k_pages.shape
   num_q_heads = queries.shape[1]
   assert num_q_heads % num_kv_heads == 0
@@ -117,6 +118,7 @@ def validate_inputs_on_runtime(
     cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
     num_seqs,  # i32[1]
 ):
+  """validate inputs on runtime"""
   check_inputs_shapes(q, k_pages, v_pages, kv_lens, page_indices, cu_q_lens, num_seqs)
   max_num_batched_tokens = q.shape[0]
   page_size = k_pages.shape[1]
@@ -148,6 +150,7 @@ def check_inputs_shapes(
     cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
     num_seqs,  # i32[1]
 ):
+  """check shapes of inputs"""
   _, num_q_heads, head_dim = q.shape
   _, _, num_kv_heads, head_dim_k = k_pages.shape
   max_num_seqs, _ = page_indices.shape
@@ -199,6 +202,7 @@ def ragged_paged_attention_kernel(
     sm_scale: float,
     mask_value: float,
 ):
+  """ragged paged-attention kernel"""
   num_q_per_blk, num_q_heads_per_blk, head_dim = q_ref.shape
   num_seqs = num_seqs_ref[0]
   _, num_kv_pages_per_blk, page_size, num_kv_heads_per_blk, _ = k_bufs.shape
@@ -522,6 +526,7 @@ def get_dtype_packing(dtype):
 
 
 def get_min_heads_per_blk(num_q_heads, num_kv_heads, q_dtype, kv_dtype):
+  """get min heads per block"""
   q_packing = get_dtype_packing(q_dtype)
   kv_packing = get_dtype_packing(kv_dtype)