code: def MBConv(
      dim_in,
      dim_out,
      *,
      downsample,
      expansion_rate=4,
      shrinkage_rate=0.25,
      dropout=0.0,
  ):
      hidden_dim = int(expansion_rate * dim_out)
      stride = 2 if downsample else 1
  
      net = nn.Sequential(
          nn.Conv2d(dim_in, hidden_dim, 1),
          nn.BatchNorm2d(hidden_dim),
          nn.GELU(),
          nn.Conv2d(
              hidden_dim,
              hidden_dim,
              3,
              stride=stride,
              padding=1,
              groups=hidden_dim,
          ),
          nn.BatchNorm2d(hidden_dim),
          nn.GELU(),
          SqueezeExcitation(hidden_dim, shrinkage_rate=shrinkage_rate),
          nn.Conv2d(hidden_dim, dim_out, 1),
          nn.BatchNorm2d(dim_out),
      )
  
      if dim_in == dim_out and not downsample:
          net = MBConvResidual(net, dropout=dropout)
  
      return net
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 4 for calling a global function with example usage: MBConv(dim_in, dim_out) and returns: net