Repository: huggingface/transformers
Pull request: Replace default split function with jnp.split() in flax models

# What does this PR do?


Fixes  (https://github.com/huggingface/transformers/issues/36854)

This issue states the error faced in RoFormer model only but few other models also faced the same issue such as 

- Albertv2
- BigBird
- Distilbert
- Electra

Attached screenshots of the error and the logs after making the changes below

changes related to the above issue are made in 

- src/transformers/models/albert/modeling_flax_albert.py
- src/transformers/models/big_bird/modeling_flax_big_bird.py
- src/transformers/models/distilbert/modeling_flax_distilbert.py
- src/transformers/models/electra/modeling_flax_electra.py
- src/transformers/models/roformer/modeling_flax_roformer.py

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://github.com/huggingface/transformers/issues/36854)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

Screenshots of models failing 

<img width="1596" alt="Screenshot 2025-03-26 at 7 07 38 PM" src="https://github.com/user-attachments/assets/bb585c23-49c0-48e1-843b-32c0cec910c2" />
<img width="1569" alt="Screenshot 2025-03-26 at 7 07 15 PM" src="https://github.com/user-attachments/assets/81745819-33bd-4521-ae2f-303315c8b5cb" />
<img width="1603" alt="Screenshot 2025-03-26 at 7 06 57 PM" src="https://github.com/user-attachments/assets/7a6c74c0-1d6c-460a-b1a9-44ba9ff87467" />
<img width="1598" alt="Screenshot 2025-03-26 at 7 06 46 PM" src="https://github.com/user-attachments/assets/0f375c06-6ce4-4379-80f6-e1edb2fcd2e8" />


Logs of test files after making the change:
[electra_split.txt](https://github.com/user-attachments/files/19468251/electra_split.txt)
[distilbert_split.txt](https://github.com/user-attachments/files/19468252/distilbert_split.txt)
[big_bird_split.txt](https://github.com/user-attachments/files/19468254/big_bird_split.txt)
[albert_split.txt](https://github.com/user-attachments/files/19468255/albert_split.txt)



## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@Rocketknight1 @gante 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @eustlb
- graph models: @clefourrier

Library:

- flax: @gante and @Rocketknight1
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @Rocketknight1
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Changed files:
filename: src/transformers/models/albert/modeling_flax_albert.py
patch: @@ -1087,7 +1087,7 @@ def __call__(
         hidden_states = outputs[0]
 
         logits = self.qa_outputs(hidden_states)
-        start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
+        start_logits, end_logits = jnp.split(logits, self.config.num_labels, axis=-1)
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)
 

filename: src/transformers/models/big_bird/modeling_flax_big_bird.py
patch: @@ -2407,7 +2407,7 @@ def __call__(
             # removing question tokens from the competition
             logits = logits - logits_mask * 1e6
 
-        start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
+        start_logits, end_logits = jnp.split(logits, self.config.num_labels, axis=-1)
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)
 

filename: src/transformers/models/distilbert/modeling_flax_distilbert.py
patch: @@ -861,7 +861,7 @@ def __call__(
 
         hidden_states = self.dropout(hidden_states, deterministic=deterministic)
         logits = self.qa_outputs(hidden_states)
-        start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
+        start_logits, end_logits = jnp.split(logits, self.config.num_labels, axis=-1)
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)
 

filename: src/transformers/models/electra/modeling_flax_electra.py
patch: @@ -1364,7 +1364,7 @@ def __call__(
         )
         hidden_states = outputs[0]
         logits = self.qa_outputs(hidden_states)
-        start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
+        start_logits, end_logits = jnp.split(logits, self.config.num_labels, axis=-1)
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)
 

filename: src/transformers/models/roformer/modeling_flax_roformer.py
patch: @@ -262,7 +262,7 @@ def __call__(
 
     @staticmethod
     def apply_rotary_position_embeddings(sinusoidal_pos, query_layer, key_layer, value_layer=None):
-        sin, cos = sinusoidal_pos.split(2, axis=-1)
+        sin, cos = jnp.split(sinusoidal_pos, 2, axis=-1)
         sin_pos = jnp.stack([sin, sin], axis=-1).reshape(sinusoidal_pos.shape)
         cos_pos = jnp.stack([cos, cos], axis=-1).reshape(sinusoidal_pos.shape)
 
@@ -1046,7 +1046,7 @@ def __call__(
         hidden_states = outputs[0]
 
         logits = self.qa_outputs(hidden_states)
-        start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
+        start_logits, end_logits = jnp.split(logits, self.config.num_labels, axis=-1)
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)