import numpy as np
import megengine
import megengine.module as M
import megengine.functional as F
import math
from . import default_init_weights

class ShuffleV2Block(M.Module):
    def __init__(self, inp, oup, mid_channels, *, ksize, stride):
        super().__init__()
        self.stride = stride
        assert stride in [1, 2]

        self.mid_channels = mid_channels
        self.ksize = ksize
        pad = ksize // 2
        self.pad = pad
        self.inp = inp

        outputs = oup - inp

        branch_main = [
            # pw
            
M.Conv2d(inp, mid_channels, 1, 1, 0, bias=False)