linting

Author: jfacevedo-google

end_to_end/tpu/eval_assert.py

@@ -22,7 +22,6 @@
 """
 
 
-
 # pylint: skip-file
 """Reads and asserts over target values"""
 from absl import app
@@ -47,7 +46,7 @@ def test_final_loss(metrics_file, target_loss, num_samples_str="10"):
   target_loss = float(target_loss)
   num_samples = int(num_samples_str)
   with open(metrics_file, "r", encoding="utf8") as _:
-    last_n_data = get_last_n_data(metrics_file, "learning/loss",num_samples)
+    last_n_data = get_last_n_data(metrics_file, "learning/loss", num_samples)
     avg_last_n_data = sum(last_n_data) / len(last_n_data)
     print(f"Mean of last {len(last_n_data)} losses is {avg_last_n_data}")
     print(f"Target loss is {target_loss}")

src/maxdiffusion/generate_wan.py

@@ -11,11 +11,7 @@
 from jax import numpy as jnp
 from safetensors import safe_open
 
-from ..modeling_flax_pytorch_utils import (
-  rename_key,
-  rename_key_and_reshape_tensor,
-  torch2jax
-)
+from ..modeling_flax_pytorch_utils import (rename_key, rename_key_and_reshape_tensor, torch2jax)
 from maxdiffusion import max_logging
 
 
@@ -36,6 +32,7 @@ class FluxParams:
   rngs: Array
   param_dtype: DTypeLike
 
+
 @dataclass
 class ModelSpec:
   params: FluxParams

src/maxdiffusion/models/flux/util.py

@@ -29,6 +29,7 @@
 
 logger = logging.get_logger(__name__)
 
+
 def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
   """
   expected_pytree: dict - a pytree that comes from initializing the model.
@@ -54,6 +55,7 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
     else:
       max_logging.log(f"key: {key} not found...")
 
+
 def torch2jax(torch_tensor: torch.Tensor) -> Array:
   is_bfloat16 = torch_tensor.dtype == torch.bfloat16
   if is_bfloat16:
@@ -67,6 +69,7 @@ def torch2jax(torch_tensor: torch.Tensor) -> Array:
   jax_array = jnp.array(numpy_value, dtype=jnp.bfloat16 if is_bfloat16 else None)
   return jax_array
 
+
 def rename_key(key):
   regex = r"\w+[.]\d+"
   pats = re.findall(regex, key)
@@ -132,7 +135,7 @@ def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dic
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
     pt_tensor = pt_tensor.transpose(2, 3, 1, 0)
     return renamed_pt_tuple_key, pt_tensor
-  
+
   # 3d conv layer
   renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 5: