text: import chainer
  # import cupy as cp
  import numpy as np
  from chainermnx.functions import allreduce
  
  
  class ChannelParallelConvolution3D(chainer.links.Convolution3D):
      def __init__(self, comm, in_channels, out_channels, *args, **kwargs):
          self.comm = comm
          self.out_channels = out_channels
          self.in_channels = in_channels
          if self.comm.size > self.in_channels:
              raise ValueError("Input channels must  be greater or equal to Ranks")
          else:
              indices = np.arange(in_channels)
              indices = indices[indices % self.comm.size == 0] + self.comm.rank
              self._channel_indices = [i for i in indices if i < self.in_channels]
              self.new_in_channels = len(self._channel_indices)
              super(ChannelParallelConvolution3D, self).__init__(self.new_in_channels, self.out_channels, *args, **kwargs)
  
      def __call__(self, x, *args, **kwargs):
          # Each process gets C/P channels
          x = x[:, self._channel_indices, :, :]
          y = super(ChannelParallelConvolution3D, self).__call__(x)
          ys = allreduce(self.comm, y)
          return ys
  
  