# AUTOGENERATED FROM: babilim/model/layers/activation.ipynb

# Cell: 0
"""doc
# babilim.model.layers.activation

> Compute an activation function.
"""

# Cell: 1
from babilim.core.annotations import RunOnlyOnce
from babilim.core.module_native import ModuleNative


# Cell: 2
class Activation(ModuleNative):
    def __init__(self, activation):
        """
        Supports the activation functions.
    
        :param activation: A string specifying the activation function to use. (Only "relu" and None supported yet.)
        """
        super().__init__()
        self.activation = activation
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        if self.activation is None:
            self.activation = self.activation
        elif self.activation == "relu":
            from torch.nn.functional import relu
            self.activation = relu
        else:
            raise NotImplementedError("Activation '{}' not implemented.".format(activation))
        
    def _call_pytorch(self, features):
        if self.activation is None:
            return features
        else:
            return self.activation(features)
    
    @RunOnlyOnce
    def _build_tf(self, features):
        if self.activation is None:
            self.activation = None
        else:
            from tensorflow.keras.layers import Activation as _Activation
            self.activation = _Activation(self.activation)
        
    def _call_tf(self, features):
        if self.activation is None:
            return features
        else:
            return self.activation(features)
