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

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

> Flatten a feature map into a linearized tensor.
"""

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


# Cell: 2
class Flatten(ModuleNative):
    def __init__(self):
        """
        Flatten a feature map into a linearized tensor.
    
        This is usefull after the convolution layers before the dense layers. The (B, W, H, C) tensor gets converted ot a (B, N) tensor.
        """
        super().__init__()
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        pass
        
    def _call_pytorch(self, features):
        return features.view(features.shape[0], -1)
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import Flatten as _Flatten
        self.flatten = _Flatten()
        self.flatten.build(features.shape)
        
    def _call_tf(self, features):
        return self.flatten(features)
