from dxtorch.dxtensor import Tensor
from .module import Module
import numpy as np


class Flatten(Module):
  def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None:
    super().__init__()
    self.start_dim = start_dim
    self.end_dim = end_dim

  def forward(self, x: Tensor) -> Tensor:
    start_dim = self.start_dim
    end_dim = self.end_dim
    if end_dim == -1:
      end_dim = x.ndim - 1
    shape = (x.shape[:self.start_dim] +
             (np.prod(x.shape[self.start_dim:self.end_dim+1]),) + x.shape[self.end_dim+1:])
    return x.reshape(shape)
