Source code for elasticai.creator.nn.binary.math_operations

from typing import cast

import torch

from elasticai.creator.base_modules.conv1d import MathOperations as Conv1dOps
from elasticai.creator.base_modules.linear import MathOperations as LinearOps
from elasticai.creator.base_modules.lstm_cell import MathOperations as LSTMOps

from .binary_quantization_function import Binarize




class MathOperations(LinearOps, Conv1dOps, LSTMOps):


    def quantize(self, a: torch.Tensor) -> torch.Tensor:
        return cast(torch.Tensor, Binarize.apply(a))




    def add(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return self.quantize(a + b)




    def mul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return self.quantize(a * b)




    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return self.quantize(torch.matmul(a, b))