einops-exts

import torch from torch import nn from einops_exts.torch import EinopsToAndFrom # Define a simple PyTorch model using EinopsToAndFrom class MyModel(nn.Module): def __init__(self): super().__init__() self.transform = EinopsToAndFrom('b h w c -> b c h w', 'b c h w -> b h w c', nn.Identity()) self.conv = nn.Conv2d(3, 64, kernel_size=3, padding=1) def forward(self, x): # Input x is expected as (batch, height, width, channels) x = self.transform(x) # Transforms to (batch, channels, height, width) for Conv2d x = self.conv(x) x = self.transform(x) # Transforms back to (batch, height, width, channels) return x # Example usage model = MyModel() input_tensor = torch.randn(1, 64, 64, 3) # Batch 1, 64x64, 3 channels output_tensor = model(input_tensor) print(f"Input shape: {input_tensor.shape}") print(f"Output shape: {output_tensor.shape}")