code: def test01():
      # conv1x1 vs linear
      from sgm.util import count_params
  
      conv = torch.nn.Conv2d(3, 32, kernel_size=1).cuda()
      print(count_params(conv))
      linear = torch.nn.Linear(3, 32).cuda()
      print(count_params(linear))
  
      print(conv.weight.shape)
  
      # use same initialization
      linear.weight = torch.nn.Parameter(
          conv.weight.squeeze(-1).squeeze(-1)
      )
      linear.bias = torch.nn.Parameter(conv.bias)
  
      print(linear.weight.shape)
  
      x = torch.randn(11, 3, 64, 64).cuda()
  
      xr = einops.rearrange(x, "b c h w -> b (h w) c").contiguous()
      print(xr.shape)
      out_linear = linear(xr)
      print(out_linear.mean(), out_linear.shape)
  
      out_conv = conv(x)
      print(out_conv.mean(), out_conv.shape)
      print("done with test01.\n")
  
file_path: generative-models/scripts/tests/attention.py
imports: import einops
  import torch
  import torch.nn.functional
  import torch.utils.benchmark
  from torch.backends.cuda import SDPBackend
  from sgm.modules.attention import BasicTransformerBlock, SpatialTransformer
label_desc: Use image node_id 4 for calling a global function with example usage: test01() without return types