Source code for elasticai.creator.base_modules.math_operations

from abc import abstractmethod
from typing import Protocol

from torch import Tensor



[docs]
class Quantize(Protocol):

[docs]
    @abstractmethod
    def quantize(self, a: Tensor) -> Tensor: ...





[docs]
class MatMul(Protocol):

[docs]
    @abstractmethod
    def matmul(self, a: Tensor, b: Tensor) -> Tensor: ...





[docs]
class Add(Protocol):

[docs]
    @abstractmethod
    def add(self, a: Tensor, b: Tensor) -> Tensor: ...





[docs]
class Mul(Protocol):

[docs]
    @abstractmethod
    def mul(self, a: Tensor, b: Tensor) -> Tensor: ...