make style

Author: bzantium

MaxText/common_types.py

@@ -81,4 +81,4 @@ class DecoderBlockType(enum.Enum):
   SIMPLE = "simple"
   SIMPLE_MLP = "simple_mlp"
   LLAMA4 = "llama4"
-  QWEN3 = "qwen3"
+  QWEN3 = "qwen3"

MaxText/layers/linears.py

@@ -84,9 +84,7 @@ def __init__(
       axis: Union[Iterable[int], int] = -1,
       weight_dtype: DType = jnp.float32,
       dtype: DType = jnp.float32,
-      kernel_init: NdInitializer = nd_dense_init(
-          1.0, "fan_in", "truncated_normal"
-      ),
+      kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"),
       kernel_axes: Tuple[Optional[str], ...] = (),
       quant: Optional[Quant] = None,
       use_bias: bool = False,
@@ -127,9 +125,7 @@ def __init__(
     # Parameter initialization
     kernel_shape = self.in_features + self.out_features
     kernel_in_axis = np.arange(len(self.axis))
-    kernel_out_axis = np.arange(
-        len(self.axis), len(self.axis) + len(self.out_features)
-    )
+    kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features))
 
     self.kernel = nnx.Param(
         self.kernel_init(
@@ -217,9 +213,7 @@ def dense_general(
     axis: Union[Iterable[int], int] = -1,
     weight_dtype: DType = jnp.float32,
     dtype: DType = jnp.float32,
-    kernel_init: NdInitializer = nd_dense_init(
-        1.0, "fan_in", "truncated_normal"
-    ),
+    kernel_init: NdInitializer = nd_dense_init(1.0, "fan_in", "truncated_normal"),
     kernel_axes: Tuple[Optional[str], ...] = (),
     quant: Optional[Quant] = None,
     use_bias: bool = False,
@@ -246,15 +240,11 @@ def dense_general(
     name: name passed to the ToLinen Module
   """
   if not (inputs_shape is not None) ^ (in_features is not None):
-    raise ValueError(
-        "Exactly one of inputs_shape or in_features must be specified."
-    )
+    raise ValueError("Exactly one of inputs_shape or in_features must be specified.")
 
   if inputs_shape is not None:
     axis = _canonicalize_tuple(axis)
-    in_features = tuple(
-        inputs_shape[ax] for ax in _normalize_axes(axis, len(inputs_shape))
-    )
+    in_features = tuple(inputs_shape[ax] for ax in _normalize_axes(axis, len(inputs_shape)))
   else:
     assert in_features is not None
   module = nnx.bridge.to_linen(
@@ -400,4 +390,3 @@ def __call__(self, inputs, decode: bool = False, deterministic: bool = False):
 
     output = checkpoint_name(output, "mlpwo")
     return output
-  

MaxText/layers/models.py

@@ -671,9 +671,7 @@ def __call__(
           inputs_shape=y.shape,
           features=cfg.vocab_size,
           weight_dtype=cfg.weight_dtype,
-          dtype=jnp.float32
-          if cfg.logits_dot_in_fp32
-          else cfg.dtype,  # for logit training stability
+          dtype=jnp.float32 if cfg.logits_dot_in_fp32 else cfg.dtype,  # for logit training stability
           kernel_axes=("embed", "vocab"),
           name="logits_dense",
           matmul_precision=self.config.matmul_precision,
@@ -809,4 +807,3 @@ def __call__(
         image_embeddings=image_embeddings,
     )
     return logits
-