text: import functools
  
  import torch.nn as nn
  
  from modules import SharedMLP, PVConv, PointNetSAModule, PointNetAModule, PointNetFPModule, PTConv
  __all__ = ['create_mlp_components', 'create_pointnet_components',
             'create_pointnet2_sa_components', 'create_pointnet2_fp_modules']
  
  
  def _linear_bn_relu(in_channels, out_channels):
      return nn.Sequential(nn.Linear(in_channels, out_channels), nn.BatchNorm1d(out_channels), nn.ReLU(True))
  
  
  def create_mlp_components(in_channels, out_channels, classifier=False, dim=2, width_multiplier=1):
      r = width_multiplier
  
      if dim == 1:
          block = _linear_bn_relu
      else:
          block = SharedMLP
      if not isinstance(out_channels, (list, tuple)):
          out_channels = [out_channels]
      if len(out_channels) == 0 or (len(out_channels) == 1 and out_channels[0] is None):
          return nn.Sequential(), in_channels, in_channels
  
      layers = []
      for oc in out_channels[:-1]:
          if oc < 1:
              layers.append(nn.Dropout(oc))
          else:
              oc = int(r * oc)
              layers.append(block(in_channels, oc))
              in_channels = oc
      if dim == 1:
          if classifier:
              layers.append(nn.Linear(in_channels, out_channels[-1]))
          else:
              layers.append(_linear_bn_relu(in_channels, int(r * out_channels[-1])))
      else:
          if classifier:
              layers.append(nn.Conv1d(in_channels, out_channels[-1], 1))
          else:
              layers.append(SharedMLP(in_channels, int(r * out_channels[-1])))
      return layers, out_channels[-1] if classifier else int(r * out_channels[-1])
  
  
  def create_pointnet_components(blocks, in_channels, with_se=False, normalize=True, eps=0,
                                 width_multiplier=1, voxel_resolution_multiplier=1):
      r, vr = width_multiplier, voxel_resolution_multiplier
  
      layers, concat_channels = [], 0
      for out_channels, num_blocks, voxel_resolution in blocks:
          out_channels = int(r * out_channels)
          if voxel_resolution is None:
              block = SharedMLP
          else:
              block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution),
                                        with_se=with_se, normalize=normalize, eps=eps)
          for _ in range(num_blocks):
              layers.append(block(in_channels, out_channels))
              in_channels = out_channels
              concat_channels += out_channels
      return layers, in_channels, concat_channels
  
  def create_pt_components(blocks, in_channels, t_layer, head, ffn_expansion=2,
          ffn_groups=1, tokenizer_tau=1, proj_tau=1, width_multiplier=1, sample_method='random'):
      r = width_multiplier
  
      layers, concat_channels = [], 0
      for out_channels, num_blocks, token_num, point_num, pos_dim in blocks:
          out_channels = int(r * out_channels)
          if token_num is None:
              block = SharedMLP
          else:
              block = functools.partial(PTConv, point_num=point_num, token_num=token_num, pos_dim=pos_dim, t_layer=t_layer, head=head, ffn_expansion=ffn_expansion,
                                        ffn_groups=ffn_groups, tokenizer_tau=tokenizer_tau, proj_tau=proj_tau, sample_method=sample_method)
          for _ in range(num_blocks):
              layers.append(block(in_channels, out_channels))
              in_channels = out_channels
              concat_channels += out_channels
      return layers, in_channels, concat_channels
  
  
  def create_pointnet2_sa_components(sa_blocks, extra_feature_channels, with_se=False, normalize=True, eps=0,
                                     width_multiplier=1, voxel_resolution_multiplier=1):
      r, vr = width_multiplier, voxel_resolution_multiplier
      in_channels = extra_feature_channels + 3
  
      sa_layers, sa_in_channels = [], []
      for conv_configs, sa_configs in sa_blocks:
          sa_in_channels.append(in_channels)
          sa_blocks = []
          if conv_configs is not None:
              out_channels, num_blocks, voxel_resolution = conv_configs
              out_channels = int(r * out_channels)
              if voxel_resolution is None:
                  block = SharedMLP
              else:
                  block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution),
                                            with_se=with_se, normalize=normalize, eps=eps)
              for _ in range(num_blocks):
                  sa_blocks.append(block(in_channels, out_channels))
                  in_channels = out_channels
              extra_feature_channels = in_channels
          num_centers, radius, num_neighbors, out_channels = sa_configs
          _out_channels = []
          for oc in out_channels:
              if isinstance(oc, (list, tuple)):
                  _out_channels.append([int(r * _oc) for _oc in oc])
              else:
                  _out_channels.append(int(r * oc))
          out_channels = _out_channels
          if num_centers is None:
              block = PointNetAModule
          else:
              block = functools.partial(PointNetSAModule, num_centers=num_centers, radius=radius,
                                        num_neighbors=num_neighbors)
          sa_blocks.append(block(in_channels=extra_feature_channels, out_channels=out_channels,
                                 include_coordinates=True))
          in_channels = extra_feature_channels = sa_blocks[-1].out_channels
          if len(sa_blocks) == 1:
              sa_layers.append(sa_blocks[0])
          else:
              sa_layers.append(nn.Sequential(*sa_blocks))
  
      return sa_layers, sa_in_channels, in_channels, 1 if num_centers is None else num_centers
  
  
  def create_pointnet2_fp_modules(fp_blocks, in_channels, sa_in_channels, with_se=False, normalize=True, eps=0,
                                  width_multiplier=1, voxel_resolution_multiplier=1):
      r, vr = width_multiplier, voxel_resolution_multiplier
  
      fp_layers = []
      for fp_idx, (fp_configs, conv_configs) in enumerate(fp_blocks):
          fp_blocks = []
          out_channels = tuple(int(r * oc) for oc in fp_configs)
          fp_blocks.append(
              PointNetFPModule(in_channels=in_channels + sa_in_channels[-1 - fp_idx], out_channels=out_channels)
          )
          in_channels = out_channels[-1]
          if conv_configs is not None:
              out_channels, num_blocks, voxel_resolution = conv_configs
              out_channels = int(r * out_channels)
              if voxel_resolution is None:
                  block = SharedMLP
              else:
                  block = functools.partial(PVConv, kernel_size=3, resolution=int(vr * voxel_resolution),
                                            with_se=with_se, normalize=normalize, eps=eps)
              for _ in range(num_blocks):
                  fp_blocks.append(block(in_channels, out_channels))
                  in_channels = out_channels
          if len(fp_blocks) == 1:
              fp_layers.append(fp_blocks[0])
          else:
              fp_layers.append(nn.Sequential(*fp_blocks))
  
      return fp_layers, in_channels
  