[BUG]einops.layer.torch.Rearrange got an unexpected keyword argument

[Wuxiaomann]

Describe the bug
When I am building ViT from monai

self.patch_embedding = PatchEmbeddingBlock(
            in_channels=in_channels,
            img_size=img_size,
            patch_size=patch_size,
            hidden_size=hidden_size,
            num_heads=num_heads,
            proj_type=proj_type,
            dropout_rate=dropout_rate,
            spatial_dims=spatial_dims,
        )

the issue happened here in monai.networks.blocks.patchembedding.py

elif self.proj_type == "perceptron":
        # for 3d: "b c (h p1) (w p2) (d p3)-> b (h w d) (p1 p2 p3 c)"
        chars = (("h", "p1"), ("w", "p2"), ("d", "p3"))[:spatial_dims]
        from_chars = "b c " + " ".join(f"({k} {v})" for k, v in chars)
        to_chars = f"b ({' '.join([c[0] for c in chars])}) ({' '.join([c[1] for c in chars])} c)"
        axes_len = {f"p{i+1}": p for i, p in enumerate(patch_size)}
        self.patch_embeddings = nn.Sequential(
            Rearrange(f"{from_chars} -> {to_chars}", **axes_len), nn.Linear(self.patch_dim, hidden_size)
        )

the error appers as
TypeError: Rearrange.init() got an unexpected keyword argument 'p1'

Reproduction steps

Expected behavior

Your platform
Version of einops, python and DL package that you used
accelerate==1.3.0
aiofiles==23.2.1
aiohttp==3.11.12
alembic==1.14.1
anyio==4.8.0
beautifulsoup4==4.13.3
bert-score==0.3.13
cachetools==5.5.1
databricks-sdk==0.44.0
deepspeed==0.16.3
docker==7.1.0
einops==0.8.1
evaluate==0.4.3
fastapi==0.115.11
Flask==3.1.0
fonttools==4.56.0
fsspec==2024.12.0
gradio==5.20.1
grpcio==1.70.0
gunicorn==23.0.0
huggingface-hub==0.28.1
imageio==2.37.0
Jinja2==3.1.5
joblib==1.4.2
lmdb==1.6.2
matplotlib==3.10.0
mlflow==2.20.1
monai==1.4.0
nibabel==5.3.2
nltk==3.9.1
numpy==1.26.4
opencv-python-headless==4.9.0.80
pandas==2.2.3
peft==0.14.0
Pillow==11.1.0
protobuf==5.29.3
psutil==6.1.1
pyarrow==18.1.0
pydantic==2.10.6
pydicom==3.0.1
python-multipart==0.0.20
scikit-learn==1.6.1
scipy==1.15.1
SimpleITK==2.4.1
tensorboard==2.19.0
timm==1.0.14
torch==2.6.0
torchvision==0.21.0
tqdm==4.67.1
transformers==4.48.3
triton==3.2.0
Unidecode==1.3.8
uvicorn==0.34.0
wandb==0.19.6

[Wuxiaomann]

And I find this issue happens when using deepspeed zero3