einops_exts_pytorch_layer_integration_with_shape_checking.py

1import torch
2from torch import nn
3from einops.layers.torch import Rearrange
4from einops_exts import check_shape, torch_wrapper
5
6# Example of using einops-exts to wrap a function for a PyTorch module
7def my_operation(x):
8    # Some operation that expects a specific shape
9    return x.mean(dim=-1)
10
11class MyModel(nn.Module):
12    def __init__(self):
13        super().__init__()
14        self.conv = nn.Conv2d(3, 16, kernel_size=3)
15        # Using Rearrange from einops directly as a layer
16        self.rearrange = Rearrange('b c h w -> b (c h w)')
17        
18    def forward(self, x):
19        # check_shape validates dimensions at runtime
20        check_shape(x, 'batch channels height width')
21        
22        x = self.conv(x)
23        x = self.rearrange(x)
24        
25        # torch_wrapper allows using einops-like string patterns with custom functions
26        # This is a conceptual example of the utility provided by exts
27        return x
28
29# Initialize and test
30model = MyModel()
31input_tensor = torch.randn(10, 3, 32, 32)
32output = model(input_tensor)
33print(f"Output shape: {output.shape}")