pre-commit

Author: sierraisland

tpu_inference/models/jax/llama4.py

@@ -3,7 +3,6 @@
 
 import jax
 import jax.numpy as jnp
-import torch
 from flax import nnx
 from flax.typing import PRNGKey
 from jax.sharding import Mesh
@@ -20,11 +19,12 @@
     SharedExpertsTransformerBlock
 from tpu_inference.logger import init_logger
 from tpu_inference.models.jax.utils.weight_utils import (
-    get_param, model_weights_generator, print_param_info, reshape_params,
-    transpose_params, convert_torch_to_jax_with_view)
+    convert_torch_to_jax_with_view, get_param, model_weights_generator,
+    print_param_info, reshape_params, transpose_params)
 
 logger = init_logger(__name__)
 
+
 class Llama4ForCausalLM(nnx.Module):
 
     def __init__(self,
@@ -121,30 +121,31 @@ def __init__(self,
                             ed_sharding=(None, 'expert'),
                             random_init=force_random_weights)
 
-            moe_ffw = MoE(dtype=dtype,
-                          num_local_experts=self.num_local_experts,
-                          apply_expert_weight_before_computation=True,
-                          hidden_size=self.hidden_size,
-                          intermediate_size_moe=self.intermediate_size_moe,
-                          hidden_act=self.hidden_act,
-                          router=router,
-                          rngs=self.rng,
-                          activation_ffw_td=('data', None),
-                          activation_ffw_ted=('data', 'expert', None),
-                          edf_sharding=('expert', None, 'model'),
-                          efd_sharding=('expert', 'model', None),
-                          random_init=force_random_weights) if is_moe_layer else None
-
-
-            dense_ffw = DenseFFW(dtype=dtype,
-                                 hidden_act=self.hidden_act,
-                                 hidden_size=self.hidden_size,
-                                 intermediate_size=self.intermediate_size_mlp,
-                                 random_init=force_random_weights,
-                                 rngs=self.rng,
-                                 df_sharding=(None, 'model'),
-                                 fd_sharding=('model', None),
-                                 activation_ffw_td=('data', None)) if not is_moe_layer else None
+            moe_ffw = MoE(
+                dtype=dtype,
+                num_local_experts=self.num_local_experts,
+                apply_expert_weight_before_computation=True,
+                hidden_size=self.hidden_size,
+                intermediate_size_moe=self.intermediate_size_moe,
+                hidden_act=self.hidden_act,
+                router=router,
+                rngs=self.rng,
+                activation_ffw_td=('data', None),
+                activation_ffw_ted=('data', 'expert', None),
+                edf_sharding=('expert', None, 'model'),
+                efd_sharding=('expert', 'model', None),
+                random_init=force_random_weights) if is_moe_layer else None
+
+            dense_ffw = DenseFFW(
+                dtype=dtype,
+                hidden_act=self.hidden_act,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size_mlp,
+                random_init=force_random_weights,
+                rngs=self.rng,
+                df_sharding=(None, 'model'),
+                fd_sharding=('model', None),
+                activation_ffw_td=('data', None)) if not is_moe_layer else None
 
             attn = Llama4Attention(
                 hidden_size=self.hidden_size,
@@ -524,8 +525,9 @@ def load_weights(self, model_for_loading: nnx.Module):
 
                     loaded_weight = convert_torch_to_jax_with_view(
                         loaded_weight, cast_type)
-                    loaded_weight = transpose_params(
-                        loaded_name, loaded_weight, self._transpose_map)
+                    loaded_weight = transpose_params(loaded_name,
+                                                     loaded_weight,
+                                                     self._transpose_map)
 
                     buffer_key = f"{mapped_name}_{'scale' if is_scale else 'qvalue'}"
                     if buffer_key not in self.expert_weights_buffer:

tpu_inference/models/jax/utils/weight_utils.py

@@ -10,10 +10,9 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
-import torch
-
 import jax
 import jax.numpy as jnp
+import torch
 from flax import nnx
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
@@ -89,7 +88,9 @@ def model_weights_generator(
                 st_file, framework, filter_regex):
             yield name, weight_tensor
 
-def convert_torch_to_jax_with_view(loaded_weight: torch.Tensor, cast_type: jnp.dtype) -> jax.Array:
+
+def convert_torch_to_jax_with_view(loaded_weight: torch.Tensor,
+                                   cast_type: jnp.dtype) -> jax.Array:
     """
     Converts a PyTorch tensor to a JAX array by reinterpreting its
     bit representation using a dtype view map.