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

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

> Reshape a tensor.
"""

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


# Cell: 2
class Reshape(ModuleNative):
    def __init__(self, output_shape):
        """
        Reshape a tensor.
    
        A tensor of shape (B, ?) where B is the batch size gets reshaped into (B, output_shape[0], output_shape[1], ...) where the batch size is kept and all other dimensions are depending on output_shape.

        :param output_shape: The shape that the tensor should have after reshaping is (batch_size,) + output_shape (meaning batch size is automatically kept).
        """
        super().__init__()
        self.output_shape = output_shape
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        self.output_shape = list(self.output_shape)
        
    def _call_pytorch(self, features):
        shape = [features.shape[0]] + self.output_shape
        return features.view(shape)
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import Reshape as _Reshape
        self.reshape = _Reshape(self.output_shape)
        
    def _call_tf(self, features):
        return self.reshape(features)
