class LinearMultiDim(nn.Linear):
    def __init__(self, in_features, out_features=None, second_dim=4, *args, **kwargs):
        in_features = [in_features, second_dim, 1] if isinstance(in_features, int) else list(in_features)
        if out_features is None:
            out_features = in_features
        out_features = [out_features, second_dim, 1] if isinstance(out_features, int) else list(out_features)
        self.in_features_multidim = in_features
        self.out_features_multidim = out_features
        super().__init__(np.array(in_features).prod(), np.array(out_features).prod())

    def forward(self, input_tensor, *args, **kwargs):
        shape = input_tensor.shape
        n_dim = len(self.in_features_multidim)
        input_tensor = input_tensor.reshape(*shape[0:-n_dim], self.in_features)
        output_tensor = super().forward(input_tensor)
        output_tensor = output_tensor.view(*shape[0:-n_dim], *self.out_features_multidim)
        return output_tensor
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py