Improve extra_repr

src/fairseq2/models/transformer/frontend.py

@@ -172,4 +172,7 @@ def extra_repr(self) -> str:
         """:meta private:"""
         s = super().extra_repr()
 
-        return f"{s}, no_scale=False" if self.scale != 1.0 else ""
+        if self.scale != 1.0:
+            s = f"{s}, no_scale=False"
+
+        return s

src/fairseq2/models/w2vbert/model.py

@@ -141,7 +141,11 @@ def layer_output_hook(
 
     def extra_repr(self) -> str:
         """:meta private:"""
-        return f"model_dim={self.model_dim}, num_bert_encoder_layers={self.num_bert_encoder_layers}, num_target_codebooks={self.num_target_codebooks}"
+        return (
+            f"model_dim={self.model_dim}, "
+            f"num_bert_encoder_layers={self.num_bert_encoder_layers}, "
+            f"num_target_codebooks={self.num_target_codebooks}"
+        )
 
 
 @dataclass

src/fairseq2/models/wav2vec2/feature_extractor.py

@@ -313,7 +313,11 @@ def _compute_seq_lens(self, num_frames: Tensor) -> Tensor:
 
     def extra_repr(self) -> str:
         """:meta private:"""
-        return f"num_fbank_channels={self.num_fbank_channels}, stride={self.stride}, sample_every_k={self.sample_every_k}"
+        return (
+            f"num_fbank_channels={self.num_fbank_channels}, "
+            f"stride={self.stride}, "
+            f"sample_every_k={self.sample_every_k}"
+        )
 
 
 class Float32LayerNorm(LayerNorm):

src/fairseq2/models/wav2vec2/masker.py

@@ -127,7 +127,12 @@ def forward(
 
     def extra_repr(self) -> str:
         """:meta private:"""
-        return f"temporal_span_len={self.temporal_span_len}, max_temporal_mask_prob={self.max_temporal_mask_prob}, spatial_span_len={self.spatial_span_len}, max_spatial_mask_prob={self.max_spatial_mask_prob}"
+        return (
+            f"temporal_span_len={self.temporal_span_len}, "
+            f"max_temporal_mask_prob={self.max_temporal_mask_prob}, "
+            f"spatial_span_len={self.spatial_span_len}, "
+            f"max_spatial_mask_prob={self.max_spatial_mask_prob}"
+        )
 
 
 def apply_temporal_mask(x: Tensor, temporal_mask: Tensor) -> Tensor:

src/fairseq2/models/wav2vec2/model.py

@@ -291,7 +291,12 @@ def _compute_logits(
 
     def extra_repr(self) -> str:
         """:meta private:"""
-        return f"model_dim={self.model_dim}, num_distractors={self.num_distractors}, logit_temp={self.logit_temp}, diversity_loss_weight={self.diversity_loss_weight}"
+        return (
+            f"model_dim={self.model_dim}, "
+            f"num_distractors={self.num_distractors}, "
+            f"logit_temp={self.logit_temp}, "
+            f"diversity_loss_weight={self.diversity_loss_weight}"
+        )
 
 
 @dataclass

src/fairseq2/nn/embedding.py

@@ -130,7 +130,9 @@ def extra_repr(self) -> str:
         s = super().extra_repr()
 
         if self.init_fn is not None:
-            s = f"{s}, init_fn={getattr(self.init_fn, '__name__', self.init_fn)}"
+            init_fn_field = getattr(self.init_fn, "__name__", self.init_fn)
+
+            s = f"{s}, init_fn={init_fn_field}"
 
         return s
 

src/fairseq2/nn/module_list.py

@@ -67,4 +67,9 @@ def drop_iter(self) -> Iterator[Module]:
 
     def extra_repr(self) -> str:
         """:meta private:"""
-        return f"drop_p={self.drop_p}" if self.drop_p > 0.0 else ""
+        s = super().extra_repr()
+
+        if self.drop_p > 0.0:
+            s = f"{s}, drop_p={self.drop_p}"
+
+        return s

src/fairseq2/nn/normalization.py

@@ -95,7 +95,11 @@ def forward(self, x: Tensor) -> Tensor:
         """
 
     def extra_repr(self) -> str:
-        return f"normalized_shape={self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}"
+        return (
+            f"normalized_shape={self.normalized_shape}, "
+            f"eps={self.eps}, "
+            f"elementwise_affine={self.elementwise_affine}"
+        )
 
 
 @final

src/fairseq2/nn/projection.py

@@ -134,7 +134,9 @@ def extra_repr(self) -> str:
         s = f"{s}, bias={self.bias is not None}"
 
         if self.init_fn is not None:
-            s = f"{s}, init_fn={getattr(self.init_fn, '__name__', self.init_fn)}"
+            init_fn_field = getattr(self.init_fn, "__name__", self.init_fn)
+
+            s = f"{s}, init_fn={init_fn_field}"
 
         return s
 

src/fairseq2/nn/transformer/__init__.py

@@ -89,3 +89,6 @@
 from fairseq2.nn.transformer.shaw_attention import (
     ShawRelativePositionSDPA as ShawRelativePositionSDPA,
 )
+from fairseq2.nn.transformer.shaw_attention import (
+    init_shaw_embedding as init_shaw_embedding,
+)

src/fairseq2/nn/transformer/decoder.py

@@ -220,8 +220,12 @@ def extra_repr(self) -> str:
         """:meta private:"""
         s = super().extra_repr()
 
-        mask_gen_name = getattr(
-            self.self_attn_mask_gen, "__name__", repr(self.self_attn_mask_gen)
+        self_attn_mask_gen_field = getattr(
+            self.self_attn_mask_gen, "__name__", self.self_attn_mask_gen
         )
 
-        return f"{s}, norm_order={self.norm_order}, self_attn_mask_gen={mask_gen_name}"
+        return (
+            f"{s}, "
+            f"norm_order={self.norm_order}, "
+            f"self_attn_mask_gen={self_attn_mask_gen_field}"
+        )

src/fairseq2/nn/transformer/ffn.py

@@ -234,4 +234,8 @@ def extra_repr(self) -> str:
         """:meta private:"""
         s = super().extra_repr()
 
-        return f"{s}, inner_dim_scale={self.inner_dim_scale}, inner_dim_to_multiple={self.inner_dim_to_multiple}"
+        return (
+            f"{s}, "
+            f"inner_dim_scale={self.inner_dim_scale}, "
+            f"inner_dim_to_multiple={self.inner_dim_to_multiple}"
+        )

src/fairseq2/nn/transformer/multihead_attention.py

@@ -603,7 +603,11 @@ def extra_repr(self) -> str:
         if self.add_zero_attn:
             s = f"{s}, add_zero_attn=True"
 
-        return f"{s}, state_factory={self.state_factory}"
+        state_factory_field = getattr(
+            self.state_factory, "__name__", self.state_factory
+        )
+
+        return f"{s}, state_factory={state_factory_field}"
 
 
 def init_qkv_projection(proj: Linear) -> None: