Fix style

vit_gpt2/modeling_flax_gpt2.py

@@ -24,7 +24,10 @@ from flax.linen.attention import dot_product_attention_weights
 from jax import lax
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_outputs import FlaxBaseModelOutputWithPastAndCrossAttentions, FlaxCausalLMOutputWithCrossAttentions
+from ...modeling_flax_outputs import (
+    FlaxBaseModelOutputWithPastAndCrossAttentions,
+    FlaxCausalLMOutputWithCrossAttentions,
+)
 from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
 from ...utils import logging
 from .configuration_gpt2 import GPT2Config
@@ -301,7 +304,9 @@ class FlaxGPT2Block(nn.Module):
         self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
 
         if self.config.add_cross_attention:
-            self.crossattention = FlaxGPT2Attention(config=self.config, dtype=self.dtype, causal=False, is_cross_attention=True)
+            self.crossattention = FlaxGPT2Attention(
+                config=self.config, dtype=self.dtype, causal=False, is_cross_attention=True
+            )
             self.ln_cross_attn = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
 
             project_encoder = getattr(self.config, "project_encoder", None)
@@ -337,7 +342,6 @@ class FlaxGPT2Block(nn.Module):
         hidden_states = attn_output + residual
 
         # Cross-Attention Block
-        cross_attn_weights = None
         if encoder_hidden_states is not None:
             # add one self-attention block for cross-attention
             if not hasattr(self, "crossattention"):
@@ -413,13 +417,16 @@ class FlaxGPT2PreTrainedModel(FlaxPreTrainedModel):
             encoder_hidden_states = jnp.zeros(input_shape + (self.config.n_embd,))
             encoder_attention_mask = attention_mask
             module_init_outputs = self.module.init(
-                rngs, input_ids, attention_mask, position_ids,
-                encoder_hidden_states, encoder_attention_mask, return_dict=False
+                rngs,
+                input_ids,
+                attention_mask,
+                position_ids,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                return_dict=False,
             )
         else:
-            module_init_outputs = self.module.init(
-                rngs, input_ids, attention_mask, position_ids, return_dict=False
-            )
+            module_init_outputs = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)
 
         return module_init_outputs["params"]
 
@@ -660,7 +667,11 @@ class FlaxGPT2Model(FlaxGPT2PreTrainedModel):
 
 
 append_call_sample_docstring(
-    FlaxGPT2Model, _TOKENIZER_FOR_DOC, _CHECKPOINT_FOR_DOC, FlaxBaseModelOutputWithPastAndCrossAttentions, _CONFIG_FOR_DOC
+    FlaxGPT2Model,
+    _TOKENIZER_FOR_DOC,
+    _CHECKPOINT_FOR_DOC,
+    FlaxBaseModelOutputWithPastAndCrossAttentions,
+    _CONFIG_FOR_DOC,
 )
 
 
@@ -718,9 +729,10 @@ class FlaxGPT2LMHeadModule(nn.Module):
             logits=lm_logits,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
-            cross_attentions=outputs.cross_attentions
+            cross_attentions=outputs.cross_attentions,
         )
 
+
 @add_start_docstrings(
     """
     The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
@@ -759,5 +771,9 @@ class FlaxGPT2LMHeadModel(FlaxGPT2PreTrainedModel):
 
 
 append_call_sample_docstring(
-    FlaxGPT2LMHeadModel, _TOKENIZER_FOR_DOC, _CHECKPOINT_FOR_DOC, FlaxCausalLMOutputWithCrossAttentions, _CONFIG_FOR_DOC
+    FlaxGPT2LMHeadModel,
+    _TOKENIZER_FOR_DOC,
+    _CHECKPOINT_FOR_DOC,
+    FlaxCausalLMOutputWithCrossAttentions,
+    _CONFIG_FOR_DOC,
 )