class_docstr: STN which learns a k-dimensional linear transformation
  
  Arguments:
      nn (torch.nn.Module) -- module which takes feat_x as input and regresses it to a global feature used to calculate the transform
      nn_feat_size -- the size of the global feature
      k -- the size of trans_x
      batch_size -- the number of examples per batch
docstr: Learns and applies a linear transformation to trans_x based on feat_x.
  feat_x and trans_x may be the same or different.
docstr_tok: ["Learns","and","applies","a","linear","transformation","to","trans_x","based","on","feat_x",".","feat_x","and","trans_x","may","be","the","same","or","different","."]
code: def forward(self, feat_x, trans_x, batch):
      """
      Learns and applies a linear transformation to trans_x based on feat_x.
      feat_x and trans_x may be the same or different.
      """
      global_feature = self.nn(feat_x, batch)
      trans = self.fc_layer(global_feature)
  
      # needed so that transform is initialized to identity
      trans = trans + self.identity.to(feat_x.device)
      trans = trans.view(-1, self.k, self.k)
      self.trans = trans
  
      # convert trans_x from (N, K) to (B, N, K) to do batched matrix multiplication
      # batch_x = trans_x.view(self.batch_size, -1, trans_x.shape[1])
      if trans_x.squeeze().dim() == 2:
          batch_x = trans_x.view(trans_x.shape[0], 1, trans_x.shape[1])
          x_transformed = torch.bmm(
              batch_x[:, :, : trans.shape[-1]], trans[batch]
          )
          if batch_x.shape[-1] > trans.shape[-1]:
              x_transformed = torch.cat(
                  [x_transformed, batch_x[:, :, trans.shape[-1] :]],
                  dim=-1,
              )
          return x_transformed.view(len(trans_x), trans_x.shape[1])
      else:
          x_transformed = torch.bmm(
              trans_x[:, :, : trans.shape[-1]], trans
          )
          if trans_x.shape[-1] > trans.shape[-1]:
              x_transformed = torch.cat(
                  [x_transformed, trans_x[:, :, trans.shape[-1] :]],
                  dim=-1,
              )
          return x_transformed
  
imports: import torch
  from torch.nn import Linear
label_desc: Use image node_id 2 for calling the BaseLinearTransformSTNkD obj's underlying member method code with example usage: obj.forward(feat_x, trans_x, batch) and returns: x_transformed, x_transformed