class_docstr: Transformer block for image-like data.
  First, project the input (aka embedding)
  and reshape to b, t, d.
  Then apply standard transformer action.
  Finally, reshape to image
class_docstr_tok: ["Transformer","block","for","image-like","data",".","First",",","project","the","input","(","aka","embedding",")","and","reshape","to","b",",","t",",","d.","Then","apply","standard","transformer","action",".","Finally",",","reshape","to","image"]
code: def forward(self, x, context=None):
      # note: if no context is given, cross-attention defaults to self-attention
      if not isinstance(context, list):
          context = [context]
      b, c, h, w = x.shape
      x_in = x
      x = self.norm(x)
      if not self.use_linear:
          x = self.proj_in(x)
      x = x.view(b, c, -1).transpose(1, 2)  # b c h w -> b (h w) c
      if self.use_linear:
          x = self.proj_in(x)
      for i, block in enumerate(self.transformer_blocks):
          x = block(x, context=context[i])
      if self.use_linear:
          x = self.proj_out(x)
      x = x.transpose(1, 2).view(b, c, h, w)  # b (h w) c -> b c h w
      if not self.use_linear:
          x = self.proj_out(x)
      return x + x_in
  
imports: import math
  from inspect import isfunction
  import torch
  import torch.nn.functional
  from apex.contrib.group_norm import GroupNorm
  from einops import rearrange
  from torch import einsum, nn
  from torch._dynamo import disable
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import checkpoint
label_desc: Use image node_id 2 for calling the SpatialTransformer obj's underlying member method code with example usage: obj.forward(x, context) and returns: unknown