from collections import OrderedDict

import torch.nn as nn
import torch.nn.functional as F

from torchmetal.modules.module import MetaModule


class MetaLinear(nn.Linear, MetaModule):
    __doc__ = nn.Linear.__doc__

    def forward(self, input, params=None):
        if params is None:
            params = OrderedDict(self.named_parameters())
        bias = params.get("bias", None)
        return F.linear(input, params["weight"], bias)


class MetaBilinear(nn.Bilinear, MetaModule):
    __doc__ = nn.Bilinear.__doc__

    def forward(self, input1, input2, params=None):
        if params is None:
            params = OrderedDict(self.named_parameters())
        bias = params.get("bias", None)
        return F.bilinear(input1, input2, params["weight"], bias)
