# stdlib
from typing import Dict
from typing import List
from typing import Tuple
from typing import Union

# relative
from ...autodp.gamma_tensor import GammaTensor
from ...autodp.phi_tensor import PhiTensor


class Layer:
    """
    Subclassed when implementing new types of layers.

    Each layer can keep track of the layer(s) feeding into it, a
    network's output :class:`Layer` instance can double as a handle to the full
    network.
    """

    first_layer: bool = False

    def forward(
        self, input: Union[PhiTensor, GammaTensor], *args: Tuple, **kwargs: Dict
    ) -> Union[PhiTensor, GammaTensor]:
        raise NotImplementedError

    def backward(
        self, pre_grad: Union[PhiTensor, GammaTensor], *args: Tuple, **kwargs: Dict
    ) -> Optional[Union[PhiTensor, GammaTensor]]:
        raise NotImplementedError

    def connect_to(self, prev_layer: "Layer") -> None:
        raise NotImplementedError

    @property
    def params(self) -> Tuple:
        """Layer parameters.

        Returns a list of numpy.array variables or expressions that
        parameterize the layer.
        Returns
        -------
        list of numpy.array variables or expressions
            A list of variables that parameterize the layer
        Notes
        -----
        For layers without any parameters, this will return an empty list.
        """
        return ()

    @property
    def grads(self) -> Tuple:
        """Get layer parameter gradients as calculated from backward()."""
        return ()

    @property
    def param_grads(self) -> List:
        """Layer parameters and corresponding gradients."""
        return list(zip(self.params, self.grads))

    def __str__(self) -> str:
        return self.__class__.__name__

@@ -50,16 +50,44 @@
 rt List%0A
+from typing import Optional%0A
 from typ

fix missing import in base.py