content: from typing import List, Optional, Tuple
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from pointnet2_ops import pointnet2_utils
  
  
  def build_shared_mlp(mlp_spec: List[int], bn: bool = True):
      layers = []
      for i in range(1, len(mlp_spec)):
          layers.append(
              nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn)
          )
          if bn:
              layers.append(nn.BatchNorm2d(mlp_spec[i]))
          layers.append(nn.ReLU(True))
  
      return nn.Sequential(*layers)
  
  
  class _PointnetSAModuleBase(nn.Module):
      def __init__(self):
          super(_PointnetSAModuleBase, self).__init__()
          self.npoint = None
          self.groupers = None
          self.mlps = None
  
      def forward(
          self, xyz: torch.Tensor, features: Optional[torch.Tensor]
      ) -> Tuple[torch.Tensor, torch.Tensor]:
          r"""
          Parameters
          ----------
          xyz : torch.Tensor
              (B, N, 3) tensor of the xyz coordinates of the features
          features : torch.Tensor
              (B, C, N) tensor of the descriptors of the the features
  
          Returns
          -------
          new_xyz : torch.Tensor
              (B, npoint, 3) tensor of the new features' xyz
          new_features : torch.Tensor
              (B,  \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
          """
  
          new_features_list = []
  
          xyz_flipped = xyz.transpose(1, 2).contiguous()
          new_xyz = (
              pointnet2_utils.gather_operation(
                  xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
              )
              .transpose(1, 2)
              .contiguous()
              if self.npoint is not None
              else None
          )
  
          for i in range(len(self.groupers)):
              new_features = self.groupers[i](
                  xyz, new_xyz, features
              )  # (B, C, npoint, nsample)
  
              new_features = self.mlps[i](new_features)  # (B, mlp[-1], npoint, nsample)
              new_features = F.max_pool2d(
                  new_features, kernel_size=[1, new_features.size(3)]
              )  # (B, mlp[-1], npoint, 1)
              new_features = new_features.squeeze(-1)  # (B, mlp[-1], npoint)
  
              new_features_list.append(new_features)
  
          return new_xyz, torch.cat(new_features_list, dim=1)
  
  
  class PointnetSAModuleMSG(_PointnetSAModuleBase):
      r"""Pointnet set abstrction layer with multiscale grouping
  
      Parameters
      ----------
      npoint : int
          Number of features
      radii : list of float32
          list of radii to group with
      nsamples : list of int32
          Number of samples in each ball query
      mlps : list of list of int32
          Spec of the pointnet before the global max_pool for each scale
      bn : bool
          Use batchnorm
      """
  
      def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True):
          # type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None
          super(PointnetSAModuleMSG, self).__init__()
  
          assert len(radii) == len(nsamples) == len(mlps)
  
          self.npoint = npoint
          self.groupers = nn.ModuleList()
          self.mlps = nn.ModuleList()
          for i in range(len(radii)):
              radius = radii[i]
              nsample = nsamples[i]
              self.groupers.append(
                  pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz)
                  if npoint is not None
                  else pointnet2_utils.GroupAll(use_xyz)
              )
              mlp_spec = mlps[i]
              if use_xyz:
                  mlp_spec[0] += 3
  
              self.mlps.append(build_shared_mlp(mlp_spec, bn))
  
  
  class PointnetSAModule(PointnetSAModuleMSG):
      r"""Pointnet set abstrction layer
  
      Parameters
      ----------
      npoint : int
          Number of features
      radius : float
          Radius of ball
      nsample : int
          Number of samples in the ball query
      mlp : list
          Spec of the pointnet before the global max_pool
      bn : bool
          Use batchnorm
      """
  
      def __init__(
          self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True
      ):
          # type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None
          super(PointnetSAModule, self).__init__(
              mlps=[mlp],
              npoint=npoint,
              radii=[radius],
              nsamples=[nsample],
              bn=bn,
              use_xyz=use_xyz,
          )
  
  
  class PointnetFPModule(nn.Module):
      r"""Propigates the features of one set to another
  
      Parameters
      ----------
      mlp : list
          Pointnet module parameters
      bn : bool
          Use batchnorm
      """
  
      def __init__(self, mlp, bn=True):
          # type: (PointnetFPModule, List[int], bool) -> None
          super(PointnetFPModule, self).__init__()
          self.mlp = build_shared_mlp(mlp, bn=bn)
  
      def forward(self, unknown, known, unknow_feats, known_feats):
          # type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
          r"""
          Parameters
          ----------
          unknown : torch.Tensor
              (B, n, 3) tensor of the xyz positions of the unknown features
          known : torch.Tensor
              (B, m, 3) tensor of the xyz positions of the known features
          unknow_feats : torch.Tensor
              (B, C1, n) tensor of the features to be propigated to
          known_feats : torch.Tensor
              (B, C2, m) tensor of features to be propigated
  
          Returns
          -------
          new_features : torch.Tensor
              (B, mlp[-1], n) tensor of the features of the unknown features
          """
  
          if known is not None:
              dist, idx = pointnet2_utils.three_nn(unknown, known)
              dist_recip = 1.0 / (dist + 1e-8)
              norm = torch.sum(dist_recip, dim=2, keepdim=True)
              weight = dist_recip / norm
  
              interpolated_feats = pointnet2_utils.three_interpolate(
                  known_feats, idx, weight
              )
          else:
              interpolated_feats = known_feats.expand(
                  *(known_feats.size()[0:2] + [unknown.size(1)])
              )
  
          if unknow_feats is not None:
              new_features = torch.cat(
                  [interpolated_feats, unknow_feats], dim=1
              )  # (B, C2 + C1, n)
          else:
              new_features = interpolated_feats
  
          new_features = new_features.unsqueeze(-1)
          new_features = self.mlp(new_features)
  
          return new_features.squeeze(-1)
  