Exception while running any model - einops.EinopsError: Error while processing rearrange-reduction pattern "b (c n) d -> b c n d".

[jchauhan]

Describe the bug

python -m examples.serving.causal-lm.llama-2-chat   --pretrained_model_name_or_path="TinyLlama/TinyLlama-1.1B-Chat-v1.0" --max_sequence_length=1024   --max_new_tokens=256  --max_compile_tokens=32 --temperature=0.6   --top_p=0.95 --top_k=50 --dtype="fp16"

   layer_outputs = block(
  File "/home/neo/research/easydel/.venv/lib/python3.10/site-packages/EasyDel/modules/llama/modelling_llama_flax.py", line 525, in __call__
    feed_forward_hidden_states = block_wise_ffn(
  File "/home/neo/research/easydel/.venv/lib/python3.10/site-packages/EasyDel/modules/flax_modelling_utils.py", line 432, in block_wise_ffn
    inputs = rearrange(inputs, 'b (c n) d -> b c n d', c=chunk_size)
  File "/home/neo/research/easydel/.venv/lib/python3.10/site-packages/einops/einops.py", line 483, in rearrange
    return reduce(cast(Tensor, tensor), pattern, reduction='rearrange', **axes_lengths)
  File "/home/neo/research/easydel/.venv/lib/python3.10/site-packages/einops/einops.py", line 420, in reduce

I have installed latest version of Easydel from github main

[erfanzar]

examples no longer work like the old version i have to replace them but the error you having right now "b (c n) d -> b c n d" is caused by block_wise_ffn what what that does is simply make MLP or FFN layer of any model chunk by chunk and then compute them you have to pass the chunk size manually but you also can turn that off by using use_scan_mlp = False in model config