Scaling issue on CPU (X86 and arm) while using models from huffing face.

I was inferring models like (BERT, Resnet, GPT2, etc)from Hugging face on CPU and found it performing was not good compared to Pytorch, so on observing found out that models are not scaling properly in FLAX case.

benchmarking results

BERT model
-------------------
cores | time in sec
1 | 38.05184221
4 | 30.66457105
8 | 29.46852469
16 | 28.98397946
32 | 28.95349479
64 | 28.77831054

### System information
- OS Platform : Ubuntu 20.04
- Flax : 0.10.0, jax: 0.4.35 , jaxlib: 0.4.35
- Python version: 3.10.12
- CPU: graviton 3

### Problem you have encountered:
scaling issue

script i am using:

cpu usage