class TorchAdd(nn.Module): """Wrapper around torch.add so that all ops can be found at build""" def __init__(self): super().__init__() self.add_func = FloatFunctional() def forward(self, x, y): return self.add_func.add(x, y)
class TorchCat2T(nn.Module): """Concatenate two tensors in channel dimension""" def __init__(self): super().__init__() self.cat_func = FloatFunctional() def forward(self, x, y): return self.cat_func.cat([x, y], dim=1)
class TorchCat(nn.Module): """Wrapper around torch.cat so that all ops can be found at build""" def __init__(self): super().__init__() self.cat_func = FloatFunctional() def forward(self, tensors: List[torch.Tensor], dim: int): return self.cat_func.cat(tensors, dim)
class TorchMultiply(nn.Module): """Wrapper around torch.mul so that all ops can be found at build""" def __init__(self): super().__init__() self.mul_func = FloatFunctional() def forward(self, x, y): return self.mul_func.mul(x, y)
class TorchMulScalar(nn.Module): """Wrapper around torch.mul so that all ops can be found at build y must be a scalar, needed for quantization """ def __init__(self): super().__init__() self.mul_func = FloatFunctional() def forward(self, x, y): return self.mul_func.mul_scalar(x, y)
class TorchAddScalar(nn.Module): """Wrapper around torch.add so that all ops can be found at build y must be a scalar, needed for quantization """ def __init__(self, scalar): super().__init__() self.add_func = FloatFunctional() self.scalar = scalar def forward(self, x): return self.add_func.add_scalar(x, self.scalar)
def __init__(self, scalar): super().__init__() self.add_func = FloatFunctional() self.scalar = scalar
def __init__(self): super().__init__() self.add_func = FloatFunctional()