class_docstr: Depth aggregator
  Args:
      aggr (str): aggregator
      input_dim (int): input dimension
      output_dim (int): output dimension
class_docstr_tok: ["Depth","aggregator","Args",":","aggr","(","str",")",":","aggregator","input_dim","(","int",")",":","input","dimension","output_dim","(","int",")",":","output","dimension"]
docstr: Forward function
  Args:
      x (torch.tensor): tensor with shape of (depth, B, N, dim)
      mask (torch.tensor): mask of x
  :return
      tensor with shape of (B, dim)
docstr_tok: ["Forward","function","Args",":","x","(","torch.tensor",")",":","tensor","with","shape","of","(","depth",",","B",",","N",",","dim",")","mask","(","torch.tensor",")",":","mask","of","x",":","return","tensor","with","shape","of","(","B",",","dim",")"]
code: def forward(self, x, mask):
      """
      Forward function
      Args:
          x (torch.tensor): tensor with shape of (depth, B, N, dim)
          mask (torch.tensor): mask of x
      :return
          tensor with shape of (B, dim)
      """
      if self.aggr is None:
          out = x[
              -1, :, 0, :
          ]  # simply takes the cls token from the last layer
      elif self.aggr == "mean":
          out = x[:, :, 0, :].mean(
              dim=0
          )  # average the cls token from the last layer
      elif self.aggr == "gated":
          mask_bool = mask.clone()
          mask_bool = mask_bool.bool()
          mask_bool = ~mask_bool
          mask_bool = mask_bool.unsqueeze(1).expand(-1, x.shape[0], -1)
          mask_bool = mask_bool.reshape(-1, mask_bool.shape[2])
  
          orig = x
          bs = x.shape[1]
          # merge batch size and depth
          x = x.view(-1, x.shape[2], x.shape[3]).permute(1, 0, 2)
          sa, _ = self.self_attn(x, x, x, key_padding_mask=mask_bool)
          scores = torch.sigmoid(self.gate_ffn(sa))  # N x bs*depth x 1
          scores = scores.permute(1, 0, 2).view(
              -1, bs, x.shape[0], 1
          )  # depth x B x N x 1
  
          # takes only the CLS
          scores = scores[:, :, 0, :]  # depth x B x 1
          orig = orig[:, :, 0, :]  # depth x B x dim
          scores = scores.permute(1, 2, 0)  # B x 1 x depth
          orig = orig.permute(1, 0, 2)  # B x depth x dim
          out = torch.matmul(scores, orig)  # B x 1 x dim
          out = out.squeeze(1)  # B x dim
  
      if self.proj is not None:
          out = self.proj(out)
      return out
  
imports: import torch
  from torch import nn
label_desc: Use image node_id 2 for calling the DepthAggregator obj's underlying member method code with example usage: obj.forward(x, mask) and returns: out