Fix Linting

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_wrappers.py

@@ -498,15 +498,17 @@ def maybe_unbox(x):
       paths_str = ""
       for path, _ in unknown_state_flat.items():
         paths_str += f"\n  - {'/'.join(map(str, path))}"
-        
+
         # Dynamically reconstruct the unknown variables
         curr = module
         for p in path[:-1]:
           if not hasattr(curr, p):
             setattr(curr, p, nnx.Module())
           curr = getattr(curr, p)
 
-      warnings.warn(f"Found unknown module paths in incoming state:{paths_str}. Intermediate modules have been reconstructed.")
+      warnings.warn(
+          f"Found unknown module paths in incoming state:{paths_str}. Intermediate modules have been reconstructed."
+      )
 
     nnx.update(module, new_state)
     _refresh_variable_trace_state(module)

src/maxtext/layers/quantizations.py

@@ -710,8 +710,10 @@ def configure_kv_quant(config):
 def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs):
   try:
     from qwix._src import flax_util
+
     parent = flax_util.get_current_module()
     from flax import nnx
+
     is_nnx = isinstance(parent, nnx.Module)
   except Exception:
     is_nnx = False
@@ -720,6 +722,7 @@ def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs):
     attr_name = f"_qwix_fp8_gpu_{op_id}"
     if not hasattr(parent, attr_name):
       from maxtext.layers import nnx_wrappers
+
       rngs = getattr(parent, "qwix_rngs", None)
       if rngs is None:
         rngs = nnx.Rngs(0)
@@ -839,10 +842,13 @@ def maybe_quantize_model(model, config):
       if config.pure_nnx:
         input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
         import jax.numpy as jnp
+
         dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
         dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
         dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
-        model = qwix.quantize_model(model, quantization_provider, dummy_tokens, dummy_positions, dummy_segment_ids, enable_dropout=False)
+        model = qwix.quantize_model(
+            model, quantization_provider, dummy_tokens, dummy_positions, dummy_segment_ids, enable_dropout=False
+        )
       else:
         model = qwix.quantize_model(model, quantization_provider)
   return model

tests/unit/quantizations_test.py

@@ -442,9 +442,10 @@ def loss_quant(model):
       # nnx.grad returns a State object which is a mapping of paths to gradients.
       # Flatten them to check for tolerance.
       from flax.nnx import traversals
+
       grads_base_flat = traversals.flatten_mapping(grads_base)
       grads_quant_flat = traversals.flatten_mapping(grads_quant)
-      
+
       # Filter for param collections to compare only parameters and not stats/buffers if any
       # Note: NNX grads structure might contain variables like 'kernel', 'bias'.
       # For simplicity we compare all matching keys.
@@ -555,7 +556,7 @@ def test_fp8_full_quantization(self):
   def test_fp8_gpu_quantization(self):
     self.quantization_config("fp8_gpu", grad_tolerance=1.5)
 
-  # @pytest.mark.gpu_only
+  @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_gpu_quantization(self):
     self.quantization_config("fp8_gpu", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)