docstr: Args:
      x: The input tensor of shape
          `(minibatch, n_pooled_feats, ..., n_ray_pts, input_dim)`.
      z: The input skip tensor of shape
          `(minibatch, n_pooled_feats, ..., n_ray_pts, skip_dim)`
          which is appended to layers whose indices are specified by `input_skips`.
  Returns:
      y: The output tensor of shape
          `(minibatch, 1, ..., n_ray_pts, input_dim)`.
code: def forward(
      self,
      x: torch.Tensor,
      z: Optional[torch.Tensor] = None,
  ):
      """
      Args:
          x: The input tensor of shape
              `(minibatch, n_pooled_feats, ..., n_ray_pts, input_dim)`.
          z: The input skip tensor of shape
              `(minibatch, n_pooled_feats, ..., n_ray_pts, skip_dim)`
              which is appended to layers whose indices are specified by `input_skips`.
      Returns:
          y: The output tensor of shape
              `(minibatch, 1, ..., n_ray_pts, input_dim)`.
      """
  
      if z is None:
          # if the skip tensor is None, we use `x` instead.
          z = x
  
      y = self.first(x)
  
      B, n_pool, n_rays, n_pts, dim = y.shape
  
      # y_p in n_pool, n_pts, B x n_rays x dim
      y_p = y.permute(1, 3, 0, 2, 4)
  
      skipi = 0
      dimh = dim
      for li, (layer_pool, layer_ray) in enumerate(
          zip(self.layers_pool, self.layers_ray)
      ):
          y_pool_attn = y_p.reshape(n_pool, n_pts * B * n_rays, dimh)
          if li in self._input_skips:
              z_skip = self.skip_linear[skipi](z)
              y_pool_attn = y_pool_attn + z_skip.permute(
                  1, 3, 0, 2, 4
              ).reshape(n_pool, n_pts * B * n_rays, dimh)
              skipi += 1
          # n_pool x B*n_rays*n_pts x dim
          y_pool_attn, pool_attn = layer_pool(
              y_pool_attn, src_key_padding_mask=None
          )
          dimh = y_pool_attn.shape[-1]
  
          y_ray_attn = (
              y_pool_attn.view(n_pool, n_pts, B * n_rays, dimh)
              .permute(1, 0, 2, 3)
              .reshape(n_pts, n_pool * B * n_rays, dimh)
          )
          # n_pts x n_pool*B*n_rays x dim
          y_ray_attn, ray_attn = layer_ray(
              y_ray_attn,
              src_key_padding_mask=None,
          )
  
          y_p = y_ray_attn.view(
              n_pts, n_pool, B * n_rays, dimh
          ).permute(1, 0, 2, 3)
  
      y = y_p.view(n_pool, n_pts, B, n_rays, dimh).permute(
          2, 0, 3, 1, 4
      )
  
      W = torch.softmax(y[..., :1], dim=1)
      y = (y * W).sum(dim=1)
      y = self.last(y)
  
      return y
  
imports: import logging
  from dataclasses import field
  from enum import Enum
  from typing import Dict, Optional, Tuple
  import torch
  from omegaconf import DictConfig
  from pytorch3d.implicitron.tools.config import Configurable, registry, ReplaceableBase, run_auto_creation
label_desc: Use image node_id 2 for calling the TransformerWithInputSkips obj's underlying member method code with example usage: obj.forward(x, z) and returns: y