oodeel 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

oodeel/methods/she.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import numpy as np
+
+from ..types import DatasetType
+from ..types import TensorType
+from ..types import Union
+from .base import OODBaseDetector
+
+
+class SHE(OODBaseDetector):
+    """
+    "Out-of-Distribution Detection based on In-Distribution Data Patterns Memorization
+    with Modern Hopfield Energy"
+    [link](https://openreview.net/forum?id=KkazG4lgKL)
+
+    This method first computes the mean of the internal layer representation of ID data
+    for each ID class. This mean is seen as the average of the ID activation patterns
+    as defined in the original paper.
+    The method then returns the maximum value of the dot product between the internal
+    layer representation of the input and the average patterns, which is a simplified
+    version of Hopfield energy as defined in the original paper.
+
+    Remarks:
+    *   An input perturbation is applied in the same way as in Mahalanobis score
+    *   The original paper only considers the penultimate layer of the neural
+    network, while we aggregate the results of multiple layers after normalizing by
+    the dimension of each vector (the activation vector for dense layers, and the
+    average pooling of the feature map for convolutional layers).
+
+    Args:
+        eps (float): magnitude for gradient based input perturbation.
+            Defaults to 0.0014.
+    """
+
+    def __init__(
+        self,
+        eps: float = 0.0014,
+    ):
+        super().__init__()
+        self.eps = eps
+        self.postproc_fns = None
+
+    def _postproc_feature_maps(self, feature_map):
+        if len(feature_map.shape) > 2:
+            feature_map = self.op.avg_pool_2d(feature_map)
+        return self.op.flatten(feature_map)
+
+    def _fit_to_dataset(
+        self,
+        fit_dataset: Union[TensorType, DatasetType],
+    ) -> None:
+        """
+        Compute the means of the input dataset in the activation space of the selected
+        layers. The means are computed for each class in the dataset.
+
+        Args:
+            fit_dataset (Union[TensorType, DatasetType]): input dataset (ID) to
+                construct the index with.
+            ood_dataset (Union[TensorType, DatasetType]): OOD dataset to tune the
+                aggregation coefficients.
+        """
+        self.postproc_fns = [
+            self._postproc_feature_maps
+            for i in range(len(self.feature_extractor.feature_layers_id))
+        ]
+
+        features, infos = self.feature_extractor.predict(
+            fit_dataset, postproc_fns=self.postproc_fns
+        )
+
+        labels = infos["labels"]
+        preds = self.op.argmax(infos["logits"], dim=-1)
+        preds = self.op.convert_to_numpy(preds)
+
+        # unique sorted classes
+        self._classes = np.sort(np.unique(self.op.convert_to_numpy(labels)))
+        labels = self.op.convert_to_numpy(labels)
+
+        self._mus = list()
+        for feature in features:
+            mus_f = list()
+            for cls in self._classes:
+                indexes = np.equal(labels, cls) & np.equal(preds, cls)
+                _features_cls = feature[indexes]
+                mus_f.append(
+                    self.op.unsqueeze(self.op.mean(_features_cls, dim=0), dim=0)
+                )
+            self._mus.append(self.op.permute(self.op.cat(mus_f), (1, 0)))
+
+    def _score_tensor(self, inputs: TensorType) -> np.ndarray:
+        """
+        Computes an OOD score for input samples "inputs" based on
+        the aggregation of neural mean discrepancies from different layers.
+
+        Args:
+            inputs: input samples to score
+
+        Returns:
+            scores
+        """
+
+        inputs_p = self._input_perturbation(inputs)
+        features, logits = self.feature_extractor.predict_tensor(
+            inputs_p, postproc_fns=self.postproc_fns
+        )
+
+        scores = self._get_she_output(features)
+
+        return -self.op.convert_to_numpy(scores)
+
+    def _get_she_output(self, features):
+        scores = None
+        for feature, mus_f in zip(features, self._mus):
+            she = self.op.matmul(self.op.squeeze(feature), mus_f) / feature.shape[1]
+            she = self.op.max(she, dim=1)
+            scores = she if scores is None else she + scores
+        return scores
+
+    def _input_perturbation(self, inputs: TensorType) -> TensorType:
+        """
+        Apply small perturbation on inputs to make the in- and out- distribution
+        samples more separable.
+
+        Args:
+            inputs (TensorType): input samples
+
+        Returns:
+            TensorType: Perturbed inputs
+        """
+
+        def __loss_fn(inputs: TensorType) -> TensorType:
+            """
+            Loss function for the input perturbation.
+
+            Args:
+                inputs (TensorType): input samples
+
+            Returns:
+                TensorType: loss value
+            """
+            # extract features
+            out_features, _ = self.feature_extractor.predict(
+                inputs, detach=False, postproc_fns=self.postproc_fns
+            )
+            # get mahalanobis score for the class maximizing it
+            she_score = self._get_she_output(out_features)
+            log_probs_f = self.op.log(she_score)
+            return self.op.mean(log_probs_f)
+
+        # compute gradient
+        gradient = self.op.gradient(__loss_fn, inputs)
+        gradient = self.op.sign(gradient)
+
+        inputs_p = inputs - self.eps * gradient
+        return inputs_p
+
+    @property
+    def requires_to_fit_dataset(self) -> bool:
+        """
+        Whether an OOD detector needs a `fit_dataset` argument in the fit function.
+
+        Returns:
+            bool: True if `fit_dataset` is required else False.
+        """
+        return True
+
+    @property
+    def requires_internal_features(self) -> bool:
+        """
+        Whether an OOD detector acts on internal model features.
+
+        Returns:
+            bool: True if the detector perform computations on an intermediate layer
+            else False.
+        """
+        return True

oodeel/methods/vim.py

@@ -61,7 +61,7 @@ class VIM(OODBaseDetector):
         pca_origin (str): either "pseudo" for using $W^{-1}b$ where $W^{-1}$ is
             the pseudo inverse of the final linear layer applied to bias term
             (as in the VIM paper), or "center" for using the mean of the data in
-            feature space. Defaults to "center".
+            feature space. Defaults to "pseudo".
     """
 
     def __init__(
@@ -86,7 +86,7 @@ class VIM(OODBaseDetector):
         """
         # extract features from fit dataset
         all_features_train, info = self.feature_extractor.predict(fit_dataset)
-        features_train = all_features_train
+        features_train = all_features_train[0]
         logits_train = info["logits"]
         features_train = self.op.flatten(features_train)
         self.feature_dim = features_train.shape[1]
@@ -101,7 +101,7 @@ class VIM(OODBaseDetector):
             # )
             W, b = self.feature_extractor.get_weights(-1)
             W, b = self.op.from_numpy(W), self.op.from_numpy(b.reshape(-1, 1))
-            _W = self.op.transpose(W) if self.backend == "tensorflow" else W
+            _W = self.op.t(W) if self.backend == "tensorflow" else W
             self.center = -self.op.reshape(self.op.matmul(self.op.pinv(_W), b), (-1,))
         else:
             raise NotImplementedError(
@@ -111,7 +111,7 @@ class VIM(OODBaseDetector):
         # compute eigvalues and eigvectors of empirical covariance matrix
         centered_features = features_train - self.center
         emp_cov = (
-            self.op.matmul(self.op.transpose(centered_features), centered_features)
+            self.op.matmul(self.op.t(centered_features), centered_features)
             / centered_features.shape[0]
         )
         eig_vals, eigen_vectors = self.op.eigh(emp_cov)
@@ -174,7 +174,7 @@ class VIM(OODBaseDetector):
         """
         # extract features
         features, logits = self.feature_extractor.predict_tensor(inputs)
-        features = self.op.flatten(features)
+        features = self.op.flatten(features[0])
         # vim score
         res_scores = self._compute_residual_score_tensor(features)
         logits = self.op.convert_to_numpy(logits)

oodeel/preprocess/__init__.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+try:
+    from .tf_preprocess import TFRandomPatchPermutation
+except ImportError:
+    pass
+
+try:
+    from .torch_preprocess import TorchRandomPatchPermutation
+except ImportError:
+    pass

oodeel/preprocess/tf_preprocess.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import numpy as np
+import tensorflow as tf
+
+from ..types import Optional
+from ..types import Tuple
+
+
+class TFRandomPatchPermutation:
+    def __init__(self, patch_size: Tuple[int] = (8, 8)):
+        """Randomly permute the patches of an image. This transformation is used in NMD
+        paper to artificially craft OOD data from ID images.
+
+        Source (NMD paper):
+            "Neural Mean Discrepancy for Efficient Out-of-Distribution Detection"
+            [link](https://arxiv.org/pdf/2104.11408.pdf)
+
+        Args:
+            patch_size (Tuple[int], optional): Patch dimensions (h, w), should be
+                divisors of the image dimensions (H, W). Defaults to (8, 8).
+        """
+        self.patch_size = patch_size
+
+    def __call__(self, tensor: tf.Tensor, seed: Optional[int] = None):
+        """Apply random patch permutation.
+
+        Args:
+            tensor (tf.Tensor): Tensor of shape [H, W, C]
+            seed (Optinal[int]): Seed number to set for the permutation if not None.
+
+        Returns:
+            tf.Tensor: Transformed tensor.
+        """
+        h, w = self.patch_size
+        H, W, C = tensor.shape
+        tensor_ = tensor
+
+        # raise warning if patch dimensions are not divisors of image dimensions
+        if H % h != 0:
+            print(
+                f"Warning! Patch height ({h}) should be a divisor of the image height"
+                + f" ({H}). Zero padding will be added to get the correct output shape."
+            )
+            tensor_ = tensor[: -(H % h)]
+        if W % w != 0:
+            print(
+                f"Warning! Patch width ({w}) should be a divisor of the image width"
+                + f" ({W}). Zero padding will be added to get the correct output shape."
+            )
+            tensor_ = tensor_[:, : -(W % w)]
+
+        # === patch permutation ===
+        # divide the batch of images into non-overlapping patches
+        # => [num_patches, h * w, C]
+        u = tf.transpose(
+            tf.reshape(tensor_, (H // h, h, W // w, w, C)), (0, 2, 1, 3, 4)
+        )
+        u = tf.reshape(u, (-1, h * w, C))
+
+        # permute the patches of each image in the batch
+        # => [num_patches, h * w, C]
+        # Note: we use numpy rng for deterministic index shuffling because
+        #       `tf.stateless_shuffle` is still experimental
+        g = np.random.default_rng(seed=seed)
+        indices = np.arange(u.shape[0])
+        g.shuffle(indices)
+        pu = tf.gather(u, indices)
+
+        # fold the permuted patches back together
+        # => [H, W, C]
+        f = tf.transpose(tf.reshape(pu, (H // h, W // w, h, w, C)), (0, 2, 1, 3, 4))
+        f = tf.reshape(f, tensor_.shape)
+        f = tf.pad(f, tf.constant([[0, H % h], [0, W % w], [0, 0]]))
+        return f

oodeel/preprocess/torch_preprocess.py

@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import torch
+import torch.nn.functional as F
+
+from ..types import Optional
+from ..types import Tuple
+
+
+class TorchRandomPatchPermutation:
+    def __init__(self, patch_size: Tuple[int] = (8, 8)):
+        """Randomly permute the patches of an image. This transformation is used in NMD
+        paper to artificially craft OOD data from ID images.
+
+        Source (NMD paper):
+            "Neural Mean Discrepancy for Efficient Out-of-Distribution Detection"
+            [link](https://arxiv.org/pdf/2104.11408.pdf)
+
+        Args:
+            patch_size (Tuple[int], optional): Patch dimensions (h, w), should be
+                divisors of the image dimensions (H, W). Defaults to (8, 8).
+        """
+        self.patch_size = patch_size
+
+    def __call__(self, tensor: torch.Tensor, seed: Optional[int] = None):
+        """Apply random patch permutation.
+
+        Args:
+            tensor (torch.Tensor): Tensor of shape [C, H, W]
+            seed (Optinal[int]): Seed number to set for the permutation if not None.
+
+        Returns:
+            torch.Tensor: Transformed tensor.
+        """
+        h, w = self.patch_size
+        H, W, _ = tensor.shape
+
+        # set generator if seed is not None
+        g = None
+        if seed is not None:
+            g = torch.Generator(device=tensor.device)
+            g.manual_seed(seed)
+
+        # raise warning if patch dimensions are not divisors of image dimensions
+        if H % h != 0:
+            print(
+                f"Warning! Patch height ({h}) should be a divisor of the image height"
+                + f" ({H}). Zero padding will be added to get the correct output shape."
+            )
+        if W % w != 0:
+            print(
+                f"Warning! Patch width ({w}) should be a divisor of the image width"
+                + f" ({W}). Zero padding will be added to get the correct output shape."
+            )
+
+        # === patch permutation ===
+        # [C, H, W] => [1, C, H, W]
+        x = tensor.unsqueeze(0)
+        # divide the batch of images into non-overlapping patches
+        # => [1, h * w, num_patches]
+        u = F.unfold(x, kernel_size=self.patch_size, stride=self.patch_size, padding=0)
+        # permute the patches of each image in the batch
+        # => [1, h * w, num_patches]
+        pu = torch.cat(
+            [b_[:, torch.randperm(b_.shape[-1], generator=g)][None, ...] for b_ in u],
+            dim=0,
+        )
+        # fold the permuted patches back together
+        # => [1, C, H, W]
+        f = F.fold(
+            pu,
+            x.shape[-2:],
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding=0,
+        )
+        return f.squeeze(0)

oodeel/utils/operator.py

@@ -49,10 +49,20 @@ class Operator(ABC):
 
     @staticmethod
     @abstractmethod
-    def max(tensor: TensorType, dim: Optional[int] = None) -> TensorType:
+    def max(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: bool = False
+    ) -> TensorType:
         """Max function"""
         raise NotImplementedError()
 
+    @staticmethod
+    @abstractmethod
+    def min(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: bool = False
+    ) -> TensorType:
+        """Min function"""
+        raise NotImplementedError()
+
     @staticmethod
     @abstractmethod
     def one_hot(tensor: TensorType, num_classes: int) -> TensorType:
@@ -138,7 +148,13 @@ class Operator(ABC):
 
     @staticmethod
     @abstractmethod
-    def transpose(tensor: TensorType) -> TensorType:
+    def t(tensor: TensorType) -> TensorType:
+        "Transpose function for tensor of rank 2"
+        raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def permute(tensor: TensorType) -> TensorType:
         "Transpose function for tensor of rank 2"
         raise NotImplementedError()
 
@@ -181,3 +197,57 @@ class Operator(ABC):
     def relu(tensor: TensorType) -> TensorType:
         "Apply relu to a tensor"
         raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def einsum(equation: str, *tensors: TensorType) -> TensorType:
+        "Computes the einsum between tensors following equation"
+        raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def tril(tensor: TensorType, diagonal: int = 0) -> TensorType:
+        "Set the upper triangle of the matrix formed by the last two dimensions of"
+        "tensor to zero"
+        raise NotImplementedError()
+
+    @staticmethod
+    def sum(tensor: TensorType, dim: Union[tuple, list, int] = None) -> TensorType:
+        "sum along dim"
+        raise NotImplementedError()
+
+    @staticmethod
+    def unsqueeze(tensor: TensorType, dim: int) -> TensorType:
+        "unsqueeze/expand_dim along dim"
+        raise NotImplementedError()
+
+    @staticmethod
+    def squeeze(tensor: TensorType, dim: int = None) -> TensorType:
+        "squeeze along dim"
+        raise NotImplementedError()
+
+    @staticmethod
+    def abs(tensor: TensorType) -> TensorType:
+        "compute absolute value"
+        raise NotImplementedError()
+
+    @staticmethod
+    def where(
+        condition: TensorType,
+        input: Union[TensorType, float],
+        other: Union[TensorType, float],
+    ) -> TensorType:
+        "Applies where function to condition"
+        raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def avg_pool_2d(tensor: TensorType) -> TensorType:
+        """Perform avg pool in 2d as in torch.nn.functional.adaptive_avg_pool2d"""
+        raise NotImplementedError()
+
+    @staticmethod
+    @abstractmethod
+    def log(tensor: TensorType) -> TensorType:
+        """Perform log"""
+        raise NotImplementedError()

oodeel/utils/tf_operator.py

@@ -65,9 +65,18 @@ class TFOperator(Operator):
         return tf.argmax(tensor, axis=dim)
 
     @staticmethod
-    def max(tensor: TensorType, dim: Optional[int] = None) -> tf.Tensor:
+    def max(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: bool = False
+    ) -> tf.Tensor:
         """Max function"""
-        return tf.reduce_max(tensor, axis=dim)
+        return tf.reduce_max(tensor, axis=dim, keepdims=keepdim)
+
+    @staticmethod
+    def min(
+        tensor: TensorType, dim: Optional[int] = None, keepdim: bool = False
+    ) -> tf.Tensor:
+        """Min function"""
+        return tf.reduce_min(tensor, axis=dim, keepdims=keepdim)
 
     @staticmethod
     def one_hot(tensor: TensorType, num_classes: int) -> tf.Tensor:
@@ -152,13 +161,18 @@ class TFOperator(Operator):
     def from_numpy(arr: np.ndarray) -> tf.Tensor:
         "Convert a NumPy array to a tensor"
         # TODO change dtype
-        return tf.constant(arr, dtype=tf.float32)
+        return tf.convert_to_tensor(arr)
 
     @staticmethod
-    def transpose(tensor: TensorType) -> tf.Tensor:
+    def t(tensor: TensorType) -> tf.Tensor:
         "Transpose function for tensor of rank 2"
         return tf.transpose(tensor)
 
+    @staticmethod
+    def permute(tensor: TensorType, dims) -> tf.Tensor:
+        "Transpose function for tensor of rank 2"
+        return tf.transpose(tensor, dims)
+
     @staticmethod
     def diag(tensor: TensorType) -> tf.Tensor:
         "Diagonal function: return the diagonal of a 2D tensor"
@@ -195,3 +209,57 @@ class TFOperator(Operator):
     def relu(tensor: TensorType) -> tf.Tensor:
         "Apply relu to a tensor"
         return tf.nn.relu(tensor)
+
+    @staticmethod
+    def einsum(equation: str, *tensors: TensorType) -> tf.Tensor:
+        "Computes the einsum between tensors following equation"
+        return tf.einsum(equation, *tensors)
+
+    @staticmethod
+    def tril(tensor: TensorType, diagonal: int = 0) -> tf.Tensor:
+        "Set the upper triangle of the matrix formed by the last two dimensions of"
+        "tensor to zero"
+        return tf.experimental.numpy.tril(tensor, k=diagonal)
+
+    @staticmethod
+    def sum(tensor: TensorType, dim: Union[tuple, list, int] = None) -> tf.Tensor:
+        "sum along dim"
+        return tf.reduce_sum(tensor, axis=dim)
+
+    @staticmethod
+    def unsqueeze(tensor: TensorType, dim: int) -> tf.Tensor:
+        "expand_dim along dim"
+        return tf.expand_dims(tensor, dim)
+
+    @staticmethod
+    def squeeze(tensor: TensorType, dim: int = None) -> tf.Tensor:
+        "expand_dim along dim"
+        return tf.squeeze(tensor, dim)
+
+    @staticmethod
+    def abs(tensor: TensorType) -> tf.Tensor:
+        "compute absolute value"
+        return tf.abs(tensor)
+
+    @staticmethod
+    def where(
+        condition: TensorType,
+        input: Union[TensorType, float],
+        other: Union[TensorType, float],
+    ) -> tf.Tensor:
+        "Applies where function to condition"
+        return tf.where(condition, input, other)
+
+    @staticmethod
+    def percentile(x, q):
+        return tfp.stats.percentile(x, q)
+
+    @staticmethod
+    def avg_pool_2d(tensor: TensorType) -> tf.Tensor:
+        """Perform avg pool in 2d as in torch.nn.functional.adaptive_avg_pool2d"""
+        return tf.reduce_mean(tensor, axis=(-3, -2))
+
+    @staticmethod
+    def log(tensor: TensorType) -> tf.Tensor:
+        """Perform log"""
+        return tf.math.log(tensor)